Python PyTorch inference_mode用法及代码示例

本文简要介绍python语言中 torch.inference_mode 的用法。

用法:

class torch.inference_mode(mode=True)

参数：

mode(bool) -标记是否启用或禁用推理模式

Context-manager 启用或禁用推理模式

InferenceMode 是一个新的上下文管理器，类似于 no_grad ，当您确定您的操作不会与 autograd 交互时使用(例如，模型训练)。在此模式下运行的代码通过禁用视图跟踪和版本计数器碰撞来获得更好的性能。

这个上下文管理器是线程本地的；它不会影响其他线程中的计算。

也起到装饰器的作用。 (确保用括号实例化。)

注意

推理模式是可以在本地启用或禁用梯度的几种机制之一，请参阅本地禁用梯度计算以获取有关它们如何比较的更多信息。

例子：

>>> import torch
>>> x = torch.ones(1, 2, 3, requires_grad=True)
>>> with torch.inference_mode():
...   y = x * x
>>> y.requires_grad
False
>>> y._version
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Inference tensors do not track version counter.
>>> @torch.inference_mode()
... def func(x):
...   return x * x
>>> out = func(x)
>>> out.requires_grad
False