登录
推荐 文章 Go 技术 课程 下载 专题 AI
首页 >  科技周边 >  人工智能

PyTorch inference_mode 与 no_grad 有什么区别

来源:17golang原创

时间:2026-09-27 07:24:48 458浏览 收藏

torch.no_grad() 和 torch.inference_mode() 都能让普通前向计算不记录反向梯度,但后者更激进:它还关闭视图追踪与版本计数器更新,并禁用 forward-mode AD,因此可能减少更多推理开销;代价是该模式创建的推理张量不能再用于被 autograd 记录的计算。

官方说明:https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.inference_mode.html

不确定输出张量是否会回到梯度计算时,优先用 no_grad;确认代码是纯数据加载、评估或生产推理,并愿意验证兼容性时,再使用 inference_mode。
要点速览
  • 两者都不会自动执行 model.eval()。
  • inference_mode 额外关闭视图追踪、版本计数器和 forward-mode AD。
  • no_grad 更宽松,适合输出仍可能参与后续 autograd 的场景。

一、共同点是关闭反向梯度记录

在两种上下文中,普通运算结果都会得到 requires_grad=False,即使输入张量原本需要梯度。它们也都是线程局部状态,可以作为上下文管理器或装饰器使用,不会直接改变其他线程。

model.eval()  # 单独切换 Dropout 和 BatchNorm 的评估行为

with torch.no_grad():
    output = model(batch)  # 不构建反向传播图

no_grad 还有一个容易忽略的例外:接受 requires_grad 参数的工厂函数不受它影响,例如在该上下文里创建 nn.Parameter,仍可能保留 requires_grad=True。

二、inference_mode 关闭更多追踪机制

inference_mode 在不记录反向梯度之外,还禁用 view tracking 和 version counter bumps。它同时关闭 forward-mode AD,因此约束比 no_grad 更强。官方建议只在确定这些操作不会与 autograd 交互时使用,并在性能收益不明确时进行基准测试。

PyTorch no_grad 与 inference_mode 对梯度、视图追踪、版本计数器和 forward-mode AD 的静态关系图
图1:静态结构图对比两种模式与反向梯度、视图追踪、版本计数器、forward-mode AD 的关系;它不是运行截图。
比较项no_gradinference_mode
反向梯度记录关闭关闭
视图追踪保留常规机制关闭
版本计数器更新保留常规机制关闭
forward-mode AD不受该 API 影响关闭
输出进入后续 autograd通常更宽松推理张量受限制

三、按张量后续用途选择

纯推理服务可以把评估状态与推理模式放在同一个清晰入口:

model.eval()  # 推理模式不会自动替代 eval()

@torch.inference_mode()
def predict(model, batch):
    logits = model(batch)  # 只用于返回预测结果
    return logits.softmax(dim=-1)

如果验证阶段还要把中间结果交给后续的梯度分析、可微后处理或 forward-mode AD,就不要只因为“看起来更快”而换成 inference_mode。这时使用 no_grad,或缩小无梯度上下文范围,更容易保持兼容。

PyTorch 训练验证、生产推理、张量复用与模式选择的静态边界图
图2:静态边界图把生产推理、验证指标、张量回到 autograd 与两种模式的选择关系放在一起,帮助判断兼容边界。

四、eval 与梯度模式解决不同问题

model.eval() 调整 Dropout、BatchNorm 等模块的运行行为;no_grad 和 inference_mode 控制自动微分记录。生产推理通常两者都要设置。只写 model.eval() 不会自动关闭梯度,只写推理上下文也不会自动切换 Dropout 与 BatchNorm。

  • 验证指标计算、输出可能再次参与梯度:优先 no_grad。
  • 稳定的纯推理或数据预处理:可尝试 inference_mode。
  • 收益不确定:用真实批量、设备和输入尺寸基准测试,不预设一定更快。

相关问题

inference_mode 会自动调用 model.eval() 吗?

不会。它控制自动微分相关机制;模型层的评估行为仍需显式调用 model.eval()。

no_grad 会影响其他线程吗?

不会。官方文档说明两种上下文都是线程局部状态,不会直接改变其他线程中的计算模式。

声明:本文转载于:17golang原创 如有侵犯,请联系study_golang@163.com删除
相关阅读
更多>
最新阅读
更多>
课程推荐
更多>