1. model.train()
如果模型中有BN层或Dropout层,model.train()是保证训练时BN层能够用到每一批数据的均值和方差,对于Dropout,model.train()是随机取一部分网络连接来训练更新参数。
2. model.eval()
model.eval()的作用是在测试时不启用 Batch Normalization 和 Dropout。在测试时,model.eval()是保证BN层能够用全部训练数据的均值和方差,即测试过程中要保证BN层的均值和方差不变;对于Dropout,model.eval()是利用到了所有网络连接,即不进行随机舍弃神经元。
3. model.train()和model.eval()使用错误示例
def test(model, test_loader):
model.eval()
...
def train(model, optimizer, epoch, train_loader, validation_loader):
model.train() # ???????????? 错误的位置
for batch_idx, (data, target) in experiment.batch_loop(iterable=train_loader):
# model.train() # 正确的位置,保证每一个batch都能进入model.train()的模式
data, target = Variable(data), Variable(target)
# Inference
output = model(data)
loss_t = F.nll_loss(output, target)
# The iconic grad-back-step trio
optimizer.zero_grad()
loss_t.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
train_loss = loss_t.item()
train_accuracy = get_correct_count(output, target) * 100.0 / len(target)
experiment.add_metric(LOSS_METRIC, train_loss)
experiment.add_metric(ACC_METRIC, train_accuracy)
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx, len(train_loader),
100. * batch_idx / len(train_loader), train_loss))
with experiment.validation():
val_loss, val_accuracy = test(model, validation_loader) # ????????????
experiment.add_metric(LOSS_METRIC, val_loss)
experiment.add_metric(ACC_METRIC, val_accuracy)
在test函数内部,将模式设置为eval,在训练过程中调用了test函数,就会进eval模式,直到下一次train函数被调用。这就导致了每一个epoch中只有一个batch使用了dropout 。
4. with torch.no_grad()
在eval模式下不会影响各层的梯度计算,即gradient计算和存储与training模式一样,只是不进行反向传播,with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。
如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation/test的结果;而with torch.no_grad()则是更进一步加速和节省gpu空间(因为不用计算和存储梯度),从而可以更快计算,也可以跑更大的batch来测试
参考
Pytorch:model.train()和model.eval()用法和区别,以及model.eval()和torch.no_grad()的区别