0
点赞
收藏
分享

微信扫一扫

model.eval()和 with torch.no_grad() 在获取中间层特征时缺一不可

谁知我新 2022-04-06 阅读 74

训练好神经网络模型后,我需要获取其中间层特征,但在通过复写模型的方法抽取特征的时候,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

举报

相关推荐

0 条评论