inference_mode¶
- class torch.autograd.grad_mode.inference_mode(mode=True)[原始碼][原始碼]¶
啟用或停用推論模式的上下文管理器。
InferenceMode 是一個上下文管理器,類似於
no_grad
,當您確定您的操作不會與 autograd 產生任何交互時(例如,模型訓練),可以使用它。在此模式下運行的程式碼通過停用視圖追蹤和版本計數器遞增來獲得更好的效能。請注意,與某些在本地啟用或停用 grad 的機制不同,進入 inference_mode 也會停用 前向模式 AD。此上下文管理器是線程本地的;它不會影響其他線程中的計算。
也可用作裝飾器。
注意
推論模式是可以本地啟用或停用梯度的幾種機制之一,有關它們如何比較的更多信息,請參閱 本地停用梯度計算。
- 參數
mode (bool 或 function) – 一個布林值,表示是否啟用或停用推論模式,或一個要使用啟用的推論模式裝飾的 Python 函數
- 範例:
>>> 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 >>> @torch.inference_mode() ... def doubler(x): ... return x * 2 >>> out = doubler(x) >>> out.requires_grad False