【踩坑】深度学习 Torch 爆显存的原因（显存不够用等问题）

简介

        在深度学习过程中，使用显卡的情况主要有两个过程：一、网络模型训练过程；二、网络模型测试过程。在这两个过程中，都可能存在爆显存或者爆内存的情况。在编程过程中，有很多同学应该都遇到这种情况，本文提供了针对这些问题的解决方案供大家参考。

        正常情况下无论是训练还是测试，显卡占用的显存资源不会大范围波动。

情况1 训练过程中爆显存

        在训练过程中，如果出现显存不够用的情况，可以先分析具体什么情况。

（1）如果瞬间爆掉显存，很大可能是因为显卡加载模型并载入训练数据时，所需要的基本显存空间不够用。解决办法是减小训练过程中一个batch的样本数，因为一个batch是一同批量输入网络模型中，一个batch中样本数越多占用的显存就越大。第二个解决办法是换个更好的显卡，越大越从容，一般做做实验使用4G或8G显存以上的显卡基本够用。

（2）如果显卡的显存占用随着训练的迭代进程逐渐升高，大概率是在代码里没有将梯度清零。

检查代码，看看有没有写

optimizer.zero_grad()    # 清空梯度

每对模型训练一次相应的模型会计算一次梯度，如果没有写这行代码，梯度会不断叠加，所以占用的显存就会逐渐增大。

情况2 测试过程中爆显存

如果是在测试过程中显存占用，可以在代码中加入

with torch.no_grad():

例如：

# 当你的代码是这么写的时候
prediction = net(testdata)

# 使用with torch.no_grad():
with torch.no_grad():
    prediction = net(testdata)

原因是当你的模型在测试数据时每次运行测试数据依旧会计算梯度并得到新的计算图，所以占用显存逐步增多，在测试过程中只需要网络计算输出结果，不需要网络计算梯度，使用

with torch.no_grad(): 

可以只计算输出结果。

当你加入了这个代码后，随着测试数据的预测进程，显存占用稳定不变。但是如果这时候你发现你的显存占用还在缓慢上升，那可能是你把预测的输出值存在类似列表等的容器中。因为网络和测试数据你都通过.cuda()移到了显卡上计算，这时候再把输出结果存在列表中，依旧是在显存内，所以随着更多的预测结果的输出，显存在缓慢增加，直到用完卡死。这时你可以把预测结果移到cpu上，利用电脑内存存储列表值。一般内存都远大于显存，而且本人发现，同样数据量下的tensor在显存上的占用比内存大很多。

with torch.no_grad():
    prediction = net(testdata)
    prediction.cpu()
    pre_list.append(prediction)

这样就可以保证你的显存不会爆掉啦。

目前发现的情况是在自己做实验时遇到的，也有可能有说错没有验证的地方，或者还有其他的情况，欢迎大家指正和补充。一起踩坑一起进步。