类原型:CLASS torch.autograd.inference_mode(mode=True)
InferenceMode是在pytorch1.10版本中引入的新功能,是一个类似于 no_grad
的新上下文管理器,该模式禁用了视图跟踪和版本计数器,所以在此模式下运行代码能够获得更好的性能,速度也会更快。
其参数表示是否启用推理模式。
以前在模型评估时,需要:
with torch.no_grad():
model.eval()
...
现在只需要更改一行代码即可:
with torch.inference_mode(mode=True):
model.eval()
...