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

最新推荐文章于 2024-12-25 16:56:51 发布

原创

最新推荐文章于 2024-12-25 16:56:51 发布 · 383 阅读

1 ·

CC 4.0 BY-SA版权

PyTorch的model.train()和model.eval()在训练和测试阶段有不同的作用。model.train()激活梯度计算，用于BN层使用每批数据的均值和方差，以及Dropout的随机失活。而model.eval()关闭梯度计算，BN层使用全局统计信息，Dropout则保持所有连接。在测试时，不加model.eval()可能会因BN层导致预测结果失真。确保在训练和测试时正确使用这两者是关键。

一、model.train()和model.eval()分别在训练和测试中都要写，它们的作用如下：

(1)、 model.train()
启用BatchNormalization和 Dropout，将BatchNormalization和Dropout置为True
(2)、 model.eval()
不启用 BatchNormalization 和 Dropout，将BatchNormalization和Dropout置为False

二、区别

如果模型中有BN层(Batch Normalization）和Dropout，需要在训练时添加model.train()，在测试时添加model.eval()。其中model.train()是保证BN层用每一批数据的均值和方差，而model.eval()是保证BN用全部训练数据的均值和方差；而对于Dropout，model.train()是随机取一部分网络连接来训练更新参数，而model.eval()是利用到了所有网络连接。

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

在做one classification的时候，训练集和测试集的样本分布是不一样的，尤其注意这一点。

使用PyTorch进行训练和测试时一定注意要把实例化的model指定train/eval，eval（）时，框架会自动把BN和DropOut固定住，不会取平均，而是用训练