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

Posted 2022-11-26 非晚非晚

• model.train()

在train模式下，dropout网络层会按照设定的参数p，设置保留激活单元的概率（保留概率=p)，并且batch norm层会继续计算数据的mean和var等参数并更新，总结如下：

• model.eval()——正常推理

在PyTorch中进行validation时，会使用model.eval()切换到测试或者称为验证模式下。model.eval()仅作用于dropout层和batchnorm层，让其与训练状态不同，相当于 dropout的参数为0 ，batchnorm也 不再更新和计算 mean和var值，而是使用训练阶段学习到的mean和var。

model.eval()不会影响各层的gradient计算行为，即gradient计算和存储与training模式一样，只是不进行反向传播。。

• with torch.no_grad()——没有梯度计算

而 with torch.no_grad()的作用是 停止对梯度的计算和存储 ，从而减少对内存的消耗，不会进行反向传播。

with torch.no_grad()并不会影响dropout和batchnorm层的行为 。

总结如下：如果不在意显存大小和计算时间的话，仅使用model.eval()已足够得到正确的validation的结果；而with torch.zero_grad()则是更进一步加速和节省gpu空间（因为不用计算和存储gradient），从而可以更快计算，也可以跑更大的batch来测试。