pytorch学习笔记-网络训练中,model.train() model.eval()的使用

1. model.train()

启用 Batch Normalization 和 Dropout。
如果模型中有BN层(Batch Normalization)和Dropout,需要在训练时添加model.train()。
model.train()作用:对BN层,保证BN层能够用到每一批数据的均值和方差,并进行计算更新;对于Dropout,model.train()是随机取一部分网络连接来训练更新参数。

2. model.eval()

不启用 Batch Normalization 和 Dropout。
如果模型中有BN层(Batch Normalization)和Dropout,在测试时添加model.eval()。
model.eval()是保证BN层直接利用之前训练阶段得到的均值和方差,即测试过程中要保证BN层的均值和方差不变;对于Dropout,model.eval()是利用到了所有网络连接,即不进行随机舍弃神经元。

3. 何时用model.eval()

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

在eval/test过程中,需要显示地让model调用eval(),此时模型会把BN和Dropout固定住,不会取平均,而是用训练好的值。

4. with torch.no_grad()

无论是train() 还是eval() 模式,各层的gradient计算和存储都在进行且完全一致,只是在eval模式下不会进行反向传播。

而with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。
若想节约算力,可在test阶段带上torch.no_grad()

示例代码:

def test(model,dataloader):
	model.eval()  # 切换到测试模式
	with torch.no_grad():  #with下内容不进行grad计算
		...

参考:Pytorch:model.train()和model.eval()用法和区别,以及model.eval()和torch.no_grad()的区别

你可能感兴趣的:(python)