pytorch 運行一段時間后出現(xiàn)GPU OOM的問題

pytorch的dataloader會將數(shù)據(jù)傳到GPU上，這個過程GPU的mem占用會逐漸增加，為了避免GPUmen被無用的數(shù)據(jù)占用，可以在每個step后用del刪除一些變量，也可以使用torch.cuda.empty_cache()釋放顯存：

del targets, input_k, input_mask
torch.cuda.empty_cache()

這時能觀察到GPU的顯存一直在動態(tài)變化。

但是上述方式不是一個根本的解決方案，因為他受到峰值的影響很大。比如某個batch的數(shù)據(jù)量明顯大于其他batch，可能模型處理該batch時顯存會不夠用，這也會導(dǎo)致OOM，雖然其他的batch都能順利執(zhí)行。

顯存的占用跟這幾個因素相關(guān)：

模型參數(shù)量

batch size

一個batch的數(shù)據(jù) size

通常我們不希望改變模型參數(shù)量，所以只能通過動態(tài)調(diào)整batch-size，使得一個batch的數(shù)據(jù) size不會導(dǎo)致顯存OOM：

ilen = int(sorted_data[start][1]['input'][0]['shape'][0])
olen = int(sorted_data[start][1]['output'][0]['shape'][0])
# if ilen = 1000 and max_length_in = 800
# then b = batchsize / 2
# and max(1, .) avoids batchsize = 0
# 太長的句子會被動態(tài)改變bsz，單獨成一個batch，否則padding的部分就太多了，數(shù)據(jù)量太大，OOM
factor = max(int(ilen / max_length_in), int(olen / max_length_out))
b = max(1, int(batch_size / (1 + factor)))
#b = batch_size
end = min(len(sorted_data), start + b)
minibatch.append(sorted_data[start:end])
if end == len(sorted_data):
    break
start = end

此外，如何選擇一個合適的batchsize也是個很重要的問題，我們可以先對所有數(shù)據(jù)按照大小（長短）排好序（降序），不進行shuffle，按照64,32,16依次嘗試bsz，如果模型在執(zhí)行第一個batch的時候沒出現(xiàn)OOM，那么以后一定也不會出現(xiàn)OOM（因為降序排列了數(shù)據(jù)，所以前面的batch的數(shù)據(jù)size最大）。

還有以下問題

pytorch increasing cuda memory OOM 問題

改了點model 的計算方式，然后就 OOM 了，調(diào)小了 batch_size,然后發(fā)現(xiàn)發(fā)現(xiàn)是模型每次迭代都會動態(tài)增長 CUDA MEMORY, 在排除了 python code 中的潛在內(nèi)存溢出問題之后，基本可以把問題定在 pytorch 的圖計算問題上了，說明每次迭代都重新生成了一張計算圖，然后都保存著在，就 OOM 了。

參考

CUDA memory continuously increases when net(images) called in every iteration

Understanding graphs and state

說是會生成多個計算圖：

loss = SomeLossFunction(out) + SomeLossFunction(out)

準(zhǔn)備用 sum來避免多次生成計算圖的問題：

loss = Variable(torch.sum(torch.cat([loss1, loss2], 0)))

然而，調(diào)著調(diào)著就好了，和報錯前的 code 沒太大差別。估計的原因是在pycharm 遠(yuǎn)程連接服務(wù)器的時候 code 的保存版本差異問題，這個也需要解決一下。

還有個多次迭代再計算梯度的問題，類似于 caffe中的iter_size,這個再仔細(xì)看看。

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。