在训练和验证模型时的model.train(),model.eval(),torch.no_grad()分别有什么作用?

本文详细解释了PyTorch中model.train()和model.eval()的作用,以及它们在启用和禁用BatchNormalization和Dropout时的差异。同时介绍了torch.no_grad()的作用,如何在验证/测试阶段优化性能和内存使用。
摘要由CSDN通过智能技术生成

引言

相信大家和我一样,与pytorch代码的结缘,都是从一些简单分类器开始的。接着我们就会发现,在不同的分类器,我们往往都会看到几句相同的代码:model.train(),model.eval(),torch.no_grad(),不免疑惑上心头,它们都是什么牛鬼蛇神呢?这时候,相信大家都毫无疑问地会发挥自己拥有一项宝贵技能——善用搜索引擎。巧了不是,俺也一样!考虑到学习是一辈子的事情,故整理了一篇笔记,记录在此。

model.train()

model.train()的作用是启用 Batch Normalization 和 Dropout。

如果模型中有BN(Batch Normalization)Dropout,需要在训练时添加model.train()。model.train()是保证BN层能够用到每一批数据的均值和方差。对于Dropoutmodel.train()是随机取一部分网络连接来训练更新参数,以使dropout层发挥它的效用。

model.eval()

model.eval()的作用是不启用 Batch Normalization 和 Dropout。

如果模型中有BN(Batch Normalization)Dropout,在测试代码前添加model.eval(),这样可以保证在测试时BN层能够用全部训练数据的均值和方差,即测试过程中要保证BN层的均值和方差不变。因为在训练模型时,BN层使用的是每一批数据的均值和方差来归一化,这样可以保证每一个batch里的数据保持统一的特征统计量(即均值和方差),从而加速模型的收敛。然而,在模型验证过程中,我们通常希望模型的预测结果与训练时保持一致,因此需要使用整个训练数据集上的均值和方差来对数据进行标准化,以保持统一的特征统计量。(看看下边这篇文章,或许会对神经网络和数据分布有一个更加宏观但清晰的认识)Batch Normalization (BN层)-----批归一化_batchnormalization层的用法-CSDN博客

对于Dropoutmodel.eval()是利用到了所有网络连接,即不进行随机舍弃神经元。

训练完train样本后,生成的模型model要用来测试样本。在model(test)之前,需要加上model.eval(),否则的话,有输入数据,即使不训练,它也会改变权值。这是model中含有BN层和Dropout所带来的的性质

model.eval()和torch.no_grad()的区别

PyTorch中进行validation/test时,会使用model.eval()切换到测试模式,在该模式下:

  1. 主要用于通知dropout层和BN层在train和validation/test模式间切换:

  • train模式下,dropout网络层会按照设定的参数p设置保留激活单元的概率(保留概率=p); BN层会继续计算数据的mean和var等参数并更新。

  • eval模式下,dropout层会让所有的激活单元都通过,而BN层会停止计算和更新meanvar,直接使用在训练阶段已经学出的meanvar值。

  • 该模式不会影响各层的gradient计算行为,即gradient计算和存储与training模式一样,只是不进行反向传播(back probagation)。

with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropoutBN层的行为。

如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation/test的结果;而with torch.no_grad()则是更进一步加速和节省gpu空间(因为不用计算和存储梯度),从而可以更快计算,也可以跑更大的batch来测试。本着能省则省的原则,建议大家加上。

总结:

综上,我们在写代码的时候,可以如下:

model.train()
# for循环的训练代码

model.eval()
with torch.no_grad():
    # 测试代码
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

不吃香菜的小趴菜

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值