PyTorch測試時候顯存爆炸問題

 

問題描述

CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 2.00 GiB total capacity; 1.13 GiB already allocated; 0 bytes free; 1.15 GiB reserved in total by PyTorch)

 

猜測：測試時候未有釋放顯卡內存，導致每次加載模型，顯卡內存都會爆炸，就很奇怪，明明測試時候只預測后面一個數據。

找了相關的文章，發現可以通過單步調試，查看顯存變化的方法。原來pytorch會把每次的梯度保存在顯存中，需要自行清空。

通過 with torch.no_grad():

語句可以不積累梯度圖。

測試結果，使用如下代碼

with torch.no_grad():

            output,hidden_state = model(x)

不再提示顯卡內存不足的情況

總結

• 梯度圖內存需要自行清理

參考：https://blog.csdn.net/dong_liuqi/article/details/119239671