Pytorch中的model. train()和model. eval()

训练完一个epoch,我们可能会生成模型来进行测试。在测试之前,需要加上model.eval(),否则的话,即使不训练,模型的权值也会改变。这是因为模型中有Batch Normalization层和Dropout层。

model.train()和model.eval()

我们知道,在pytorch中,模型有两种模式可以设置,一个是train模式、另一个是eval模式。

model.train():的作用是启用 Batch Normalization 和 Dropout。在train模式,Dropout层会按照设定的参数p设置保留激活单元的概率,如keep_prob=0.8,Batch Normalization层会继续计算数据的mean和var并进行更新。

model.eval():的作用是不启用 Batch Normalization 和 Dropout。在eval模式下,Dropout层会让所有的激活单元都通过,而Batch Normalization层会停止计算和更新mean和var,直接使用在训练阶段已经学出的mean和var值。 在使用model.eval()时就是将模型切换到测试模式,在这里,模型就不会像在训练模式下一样去更新权重。但是需要注意的是model.eval()不会影响各层的梯度计算行为,即会和训练模式一样进行梯度计算和存储,只是不进行反向传播。

# 针对BN层:
model.train() # 是保证BN层用每一批数据的均值和方差,即针对每个mini-batch的 ;
model.eval()  # 是保证BN用全部训练数据的均值和方差,即针对单张图片的;
# 针对Dropout层
model.train() # 随机取一部分网络连接来训练更新参数;
model.eval()  # 利用到了所有网络连接;

model.eval()和torch.no_grad()

在讲model.eval()时,其实还会提到torch.no_grad()。

torch.no_grad():用于停止autograd的计算,能起到加速和节省显存的作用,但是不会影响Dropout层和Batch Normalization层的行为。

如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation的结果;而

with torch.zero_grad():则是更进一步加速和节省gpu空间。因为不用计算和存储梯度,从而可以计算得更快,也可以使用更大的batch来运行模型。

详细分析,可以看一下:【PyTorch】搞定网络训练中的model.train()和model.eval()模式 - 知乎

猜你喜欢

转载自blog.csdn.net/ytusdc/article/details/128523707
今日推荐