训练好神经网络模型后,我需要获取其中间层特征,但在通过复写模型的方法抽取特征的时候,vscode多次崩溃,并提醒我buy a new RAM!
后来经排查发现,我只进行了 model.eval()操作,而没有在抽取中间层的时候加上 with torch.no_grad()。
model.eval()仅作用于dropout层和batchnorm层,让其与训练状态不同,相当于dropout的参数为0,batchnorm也不再更新和计算mean和var值,而是使用训练阶段学习到的mean和var。
而 with torch.no_grad()的作用是停止对梯度的计算和存储,从而减少对内存的消耗。
因此,在电脑硬件资源有限的情况下,model.eval()和with torch.no_grad()缺一不可!
参考:
(2条消息) pytorch中model eval和torch no grad()的区别_江前云后的博客-CSDN博客_torch.nograd