在GPU上训练神经网络时如何解决OOM问题?

时间:2019-02-15 07:59:01

标签: python tensorflow neural-network

我正在使用TensorFlow在NVIDIA GEFORCE GTX 1080 Ti上训练神经网络。我的CPU有32 GB RAM。 该网络只有1个隐藏层,其权重分别为[16788,10000]和[10000,1478]。最小批量= 256。

a.b.c.d

成功完成3个时期后,会弹出以下错误消息:

for epo in range(self.epochs):
    loss_sum = 0        
    for mini_count in (range(len(mini_batches_x1_train))):
        batch_x1 = tf.one_hot(mini_batches_x1_train[mini_count], self.voc_size).eval()
        batch_x2 = mini_batches_x2_train[mini_count][0].toarray() 
        batch_x = np.concatenate((batch_x1, batch_x2), axis = 1)

        batch_y = mini_batches_y_train[mini_count]
        batch_y = np.array(batch_y).reshape(len(batch_y),1)  #for tf.nce_loss()
        _, mini_loss = sess.run([optimizer,loss], feed_dict={X: batch_x, Y: batch_y})
        loss_sum += mini_loss

    avg_loss = loss_sum/len(mini_batches_x1_train)
    print("\nEpoch", epo+1, " completed at ",time.ctime(time.time()), " | Epoch Loss = ", avg_loss)

更多的纪元需要更多的内存吗?我在哪里可能错了?请提出解决此问题的可能解决方案。

1 个答案:

答案 0 :(得分:0)

在GPU上进行训练时,Tensorflow会尝试尽可能多地将其加载到GPU VRAM中,以使其速度更快且不会因频繁访问GPU中的RAM数据而成为瓶颈。 因此,如果网络不适合VRAM,则基本上有多种选择。

  1. 最广泛使用的选项是减小批量大小。尝试使用128或64(建议将批量大小的幂乘以2,请参见this SO question)。某些最先进的架构是如此之大,即使在GTX 1080Ti上,批次大小也只能= 1。
  2. 改变架构,使其更小。
  3. 在CPU上运行模型(出于性能考虑,您可能不希望这样做)。
  4. 将优化器更改为一些使用较低内存的优化器。例如。用于获得SOTA性能的ADAM优化器将其自己的变量添加到每一层,因此,与某些简单的优化器相比,整个模型消耗更多的内存。尽管通常不想更改优化器,而只是减小批处理大小。