引言
相信大家和我一样,与pytorch代码的结缘,都是从一些简单分类器开始的。接着我们就会发现,在不同的分类器,我们往往都会看到几句相同的代码:model.train(),model.eval(),torch.no_grad(),不免疑惑上心头,它们都是什么牛鬼蛇神呢?这时候,相信大家都毫无疑问地会发挥自己拥有一项宝贵技能——善用搜索引擎。巧了不是,俺也一样!考虑到学习是一辈子的事情,故整理了一篇笔记,记录在此。
model.train()
model.train()
的作用是启用 Batch Normalization 和 Dropout。
如果模型中有BN
层(Batch Normalization)
和Dropout
,需要在训练时添加model.train()。model.train()
是保证BN层能够用到每一批数据的均值和方差。对于Dropout
,model.train()
是随机取一部分网络连接来训练更新参数,以使dropout层发挥它的效用。
model.eval()
model.eval()的作用是不启用 Batch Normalization 和 Dropout。
如果模型中有BN
层(Batch Normalization)
和Dropout
,在测试代码前添加model.eval(),这样可以
保证在测试时BN层能够用全部训练数据的均值和方差,即测试过程中要保证BN层的均值和方差不变。因为在训练模型时,BN层使用的是每一批数据的均值和方差来归一化,这样可以保证每一个batch里的数据保持统一的特征统计量(即均值和方差),从而加速模型的收敛。然而,在模型验证过程中,我们通常希望模型的预测结果与训练时保持一致,因此需要使用整个训练数据集上的均值和方差来对数据进行标准化,以保持统一的特征统计量。(看看下边这篇文章,或许会对神经网络和数据分布有一个更加宏观但清晰的认识)Batch Normalization (BN层)-----批归一化_batchnormalization层的用法-CSDN博客
对于Dropout
,model.eval()
是利用到了所有网络连接,即不进行随机舍弃神经元。
训练完train
样本后,生成的模型model
要用来测试样本。在model(test)
之前,需要加上model.eval()
,否则的话,有输入数据,即使不训练,它也会改变权值。这是model
中含有BN层和Dropout
所带来的的性质
model.eval(
)和torch.no_grad()
的区别
在PyTorch
中进行validation/test
时,会使用model.eval()
切换到测试模式,在该模式下:
-
主要用于通知dropout层和BN层在train和validation/test模式间切换:
-
在
train
模式下,dropout
网络层会按照设定的参数p设置保留激活单元的概率(保留概率=p); BN层会继续计算数据的mean和var等参数并更新。 -
在
eval
模式下,dropout
层会让所有的激活单元都通过,而BN
层会停止计算和更新mean
和var
,直接使用在训练阶段已经学出的mean
和var
值。 -
该模式不会影响各层的
gradient
计算行为,即gradient
计算和存储与training
模式一样,只是不进行反向传播(back probagation)。
而with torch.no_grad()
则主要是用于停止autograd
模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU
算力和显存,但是并不会影响dropout
和BN
层的行为。
如果不在意显存大小和计算时间的话,仅仅使用model.eval()
已足够得到正确的validation/test
的结果;而with torch.no_grad()
则是更进一步加速和节省gpu
空间(因为不用计算和存储梯度),从而可以更快计算,也可以跑更大的batch
来测试。本着能省则省的原则,建议大家加上。
总结:
综上,我们在写代码的时候,可以如下:
model.train()
# for循环的训练代码
model.eval()
with torch.no_grad():
# 测试代码