评估 pytorch 模型:`with torch.no_grad` vs `model.eval()`

Evaluating pytorch models: `with torch.no_grad` vs `model.eval()`

当我想评估我的模型在验证集上的性能时,首选使用with torch.no_grad:还是model.eval()

长话短说:

Use both。他们做不同的事情,有不同的范围。

  • with torch.no_grad - 在 autograd.
  • 中禁用渐变跟踪
  • model.eval() 改变了它被调用的模块的 forward() 行为
  • 例如,它禁用了 dropout 并让 batch norm 使用整个人口统计数据

with torch.no_grad

torch.autograd.no_grad documentation 说:

Context-manager that disabled [sic] gradient calculation.

Disabling gradient calculation is useful for inference, when you are sure that you will not call Tensor.backward(). It will reduce memory consumption for computations that would otherwise have requires_grad=True. In this mode, the result of every computation will have requires_grad=False, even when the inputs have requires_grad=True.

model.eval()

nn.Module.eval documentation 说:

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.


The creator of pytorch said the documentation should be updated to suggest the usage of both, and I raised the pull request.

with torch.no_grad: 禁用 backward 过程的梯度计算。由于这些计算在推理过程中是不必要的,并且会增加不平凡的计算开销,因此如果评估模型的速度,则必须使用此上下文。但是不会影响结果。

model.eval() 确保 forward 推理过程中正确定义在训练与推理中表现不同的行为(例如 Dropout 和 BatchNorm)。因此,如果您的模型包含此类模块,则必须启用它。

出于上述原因,在推理过程中同时使用两者是一种很好的做法。

如果您正在阅读此 post 是因为您遇到过 RuntimeError: CUDA out of memory,那么 with torch.no grad(): 可能有助于节省记忆。仅使用 model.eval() 不太可能有助于解决 OOM 错误。

原因是torch.no grad() 完全禁用了 autograd(你不能再反向传播),减少了内存消耗并加快了计算速度。

但是,您仍然可以在使用 model.eval() 时呼叫 gardients。就个人而言,我觉得这个设计决定很有趣。那么,.eval()的目的是什么?它的主要功能似乎是在评估期间停用 Dropout。

总而言之,如果您使用 torch.no grad(),则不会保存任何中间张量,并且您可能会增加推理中的批量大小。