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

Question

我正在使用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)

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

Answer 1

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

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