Python PyTorch inference_mode用法及代碼示例

本文簡要介紹python語言中 torch.autograd.inference_mode 的用法。

用法:

class torch.autograd.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