推理模式¶
- class torch.autograd.grad_mode.inference_mode(mode=True)[源代码]¶
上下文管理器,用于启用或禁用推理模式。
InferenceMode 是一个新的上下文管理器,类似于
no_grad, 用于在你确定操作不会与自动求导(例如,模型训练)有交互时使用。在此模式下运行的代码通过禁用视图跟踪和版本计数器更新来获得更好的性能。请注意, 与一些局部启用或禁用梯度的机制不同,进入 inference_mode 还会禁用 前向模式 AD。此上下文管理器是线程本地的;它不会影响其他线程中的计算。
也可以作为装饰器使用。
注意
推理模式是几种可以在本地启用或禁用梯度的机制之一,有关它们如何比较的更多信息,请参见本地禁用梯度计算。
- Parameters
模式 (布尔值 或 函数) – 一个布尔标志,用于启用或禁用推理模式,或者是一个Python函数,用于在启用推理模式的情况下进行装饰
- Example::
>>> 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 "
", line 1, in RuntimeError: 推理张量不跟踪版本计数器。 >>> @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