pytorch之model.cuda()、model.train()和model.eval()

2022-10-18 09:16:22

model.cuda()

在pytorch中,即使是有GPU的机器,它也不会自动使用GPU,而是需要在程序中显示指定。调用model.cuda(),可以将模型加载到GPU上去。

model.train()

启用 BatchNormalization 和 Dropout

model.eval()

不启用 BatchNormalization 和 Dropout

model.train()&&model.eval()

PyTorch进行训练和测试时一定注意要把实例化的model指定train/eval
eval()时,框架会自动把 BN 和 DropOut 固定住,不会取平均,而是用训练好的值,不然的话,一旦test的batch_size过小,很容易就会被BN层导致生成图片颜色失真极大!
训练完train样本后,生成的模型model要用来测试样本。在model(test)之前,需要加上model.eval(),否则的话,有输入数据,即使不训练,它也会改变权值。这是model中含有batch normalization层所带来的的性质。

  • 作者:ZwaterZ
  • 原文链接:https://blog.csdn.net/ZwaterZ/article/details/125651755
    更新时间:2022-10-18 09:16:22