在GPU上訓練神經網絡時如何解決OOM問題？

Question

我正在使用NVIDIA GEFORCE GTX 1080 Ti上的TensorFlow訓練神經網絡。 我的CPU有32 GB RAM。 該網絡只有1個隱藏層，其權重分別為[16788,10000]和[10000，1478]。 迷你批次大小= 256。

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)

成功完成3個時期后，會彈出以下錯誤：

ResourceExhaustedError: OOM when allocating tensor with shape[16788,10000] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

更多的紀元需要更多的內存嗎？ 我在哪里可能錯了？ 請提出可能的解決方案以克服這一問題。

Answer 1

在GPU上進行訓練時，Tensorflow會嘗試盡可能多地將其加載到GPU VRAM中，以使其速度更快且不會因過於頻繁地訪問GPU中的RAM數據而成為瓶頸。 因此，如果網絡不適合VRAM，則基本上有多種選擇。

1. 最廣泛使用的選項是減小批次大小。 嘗試使用128或64（建議將批大小的冪保持為2，請參見此SO問題 ）。 一些最先進的架構是如此龐大，即使在GTX 1080Ti上，批次大小也只能= 1。
2. 更改體系結構並使其更小。
3. 在CPU上運行模型（由於性能，您可能不希望這樣做）。
4. 將優化器更改為一些使用較低內存的優化器。 例如，用於獲得SOTA性能的ADAM優化器將其自己的變量添加到每一層，因此，與某些簡單的優化器相比，整個模型消耗的內存更多。 盡管通常不想更改優化器，而只是減小批處理大小。