[已解決][PyTorch] RuntimeError: CUDA out of memory. Tried to allocate 2.0 GiB

Last Updated on 2021-07-11 by Clay

今天我想要紀錄一個常見的問題，但是其解決方法其實非常少見。簡單來講，我的報錯顯示如下：

RuntimeError: CUDA out of memory. Tried to allocate 2.0 GiB.

這個報錯其實非常單純，那就是 GPU 的『記憶體』不夠了，導致我們想要在 GPU 內執行的訓練資料不夠存放，導致程式意外中止。

是的，使用 nvidia-smi 看到的記憶體容量是 GPU 的記憶體；而使用 htop 所查看到的記憶體則是一般電腦內正常存放執行程序的記憶體，兩者是不一樣的。

解決方法

通常如果是在訓練階段遇到這個問題，通常都是 Batch Size 過大的問題。試想：一次丟給 GPU 很大很大的資料量，是不是容易讓記憶體爆掉呢？反之，若是一次丟的資料小一點，然後訓練完後就清除掉，換下一批資料進來，是不是就不會讓 GPU 爆掉了呢？

所以如果是訓練階段，調小 Batch Size 是可以考慮的方法。

不過我今天遇到的問題稍微複雜點，我訓練了一個模型，但是我接下來讀取這個模型只希望將其取用第一層的輸出 —— 在我重複這個動作一陣子之後，我 GPU 的記憶體就爆掉了！

這是因為，雖然我將模型調整為 eval() 模式，但是資料的梯度仍然在累積，而且又是累積在 GPU 中。久而久之，GPU 的記憶體就超過上限了。

解決方法是：在我們要使用模型產出的時候，將其放在 with torch.no_grad() 底下，這樣一來，模型就不會累積梯度。