PyTorch C++ 梯度模式守卫实战指南：NoGradGuard 与 InferenceMode 的机制、性能差异与迁移路径

2026-09-07 14:25:11作者：柏廷章Berta

本指南围绕 PyTorch C++ API 中用于控制自动梯度（autograd）计算行为的 RAII 作用域守卫展开，系统讲解 torch::NoGradGuard 与 c10::InferenceMode 的用法、底层 TLS 状态与 DispatchKey 实现、推理张量（inference tensor）的约束规则，以及从遗留的 AutoNonVariableTypeMode 安全迁移的完整方案。读完本文，你将能够在纯推理、推理+后处理混合以及自定义 autograd kernel 三种场景下，准确选出并正确使用对应的梯度模式守卫，理解为何"只看推理不看训练"的代码也能在不损失正确性的前提下获得更优性能。

梯度模式：PyTorch C++ 控制梯度计算的 RAII 机制

在 PyTorch 的 C++ 前向接口中，autograd 系统是否记录 grad_fn、是否进行视图追踪（view tracking）、是否更新版本计数器（version counter）等行为，由一组线程局部（thread local）的状态标志控制。这组标志被集中封装在 c10::AutogradState 中，包含 grad_mode、inference_mode、fw_grad_mode 与 multithreading_enabled 等位域（见 c10/core/AutogradState.h）。

为了以"作用域"为单位干净地切换这些状态，PyTorch 提供了基于 RAII 的守卫类型：在构造时保存当前状态并施加新的模式，在析构时自动恢复先前的状态。官方 C++ API 文档明确支持的两种核心守卫是：

NoGradGuard：在其作用域内禁用梯度计算；
InferenceMode（即 c10::InferenceMode）：为确定不会与 autograd 产生任何交互的推理负载提供更激进的优化模式。

两者都遵循"构造即生效、析构即还原"的 RAII 语义，且状态为线程局部（thread local）——相关实现注释见 c10/core/InferenceMode.h。

NoGradGuard：作用域内禁用梯度计算

定义与语义

NoGradGuard 是一个 RAII 守卫，用于在其作用域内禁用梯度计算：

torch::NoGradGuard no_grad;      // 构造：禁用梯度计算
// ... 作用域内的所有算子不再记录 grad_fn ...
~NoGradGuard();                  // 析构：恢复进入前的梯度模式

从实现上看，它并非独立的类，而是直接继承 c10::AutoGradMode，并通过把 enabled 置为 false 来关闭梯度模式：

// c10/core/GradMode.h
struct C10_API NoGradGuard : public AutoGradMode {
  NoGradGuard() : AutoGradMode(/*enabled=*/false) {}
};

同时，在 ATen 与 torch::autograd 命名空间中，它被以别名形式导出（见 aten/src/ATen/core/grad_mode.h 与 torch/csrc/autograd/grad_mode.h），因此在文档示例中可以直接写作 torch::NoGradGuard。

基本用法

{
    torch::NoGradGuard no_grad;
    // No gradients computed in this scope
    auto result = model->forward(input);
}   // 离开作用域后，梯度模式恢复为之前的状态

要点：

守卫只需要声明为局部变量，析构由作用域退出自动触发；
通常用于只在个别代码段临时关闭梯度（例如推理代码与训练逻辑交织）的场景；
它只影响"是否记录 grad_fn"，不影响张量本身的创建方式——在 NoGradGuard 内创建的张量与普通张量没有区别，离开作用域后依然可以自由地被用于求梯度、原地修改等操作。

InferenceMode：面向纯推理工作负载的高性能守卫

定位与设计动机

c10::InferenceMode 是一个与 NoGradMode 语义类似、但专为确定不会与 autograd 交互的代码（例如模型推理）设计的 RAII 守卫。文档给出的核心定位是：相比 NoGradMode，InferenceMode 下运行的代码通过跳过视图追踪（view tracking）与版本计数器自增（version counter bumps）等 autograd 相关工作获得更好的性能；代价是，在其内部创建的张量在与 autograd 系统交互时会受到更多限制。

c10::InferenceMode guard;  // 等价于 c10::InferenceMode guard(true);

构造函数原型为 explicit InferenceMode(bool enabled = true)，即既可以启用也可以显式禁用。

源码级实现：TLS 状态与 DispatchKeySet 的双重切换

从 c10/core/InferenceMode.h 的实现可以看出，InferenceMode 的构造并不仅仅是关闭 GradMode，而是同时做三件事：

同步设置 autograd TLS 状态：
```
AutogradState::set_tls_state(AutogradState(
    /* grad_mode */ !enabled,
    /* inference_mode */ enabled,
    /* fw_grad_mode */ !enabled,
    /* multithreading_enabled*/ !enabled));
```
也就是说，开启 InferenceMode 的同时会把 grad_mode 与正向梯度模式 fw_grad_mode 一并关闭。这正是源码注释第 3 点强调的："InferenceMode is a faster and more restrictive version of NoGradGuard"——所有基于 GradMode::is_enabled() 的运行时检查对 InferenceMode 同样成立。
从本地 DispatchKeySet 中移除 ADInplaceOrView key（启用时）：
```
DispatchKeySet included = enabled
    ? prev_keyset.included_.remove(c10::DispatchKey::ADInplaceOrView)
    : prev_keyset.included_.add(c10::DispatchKey::ADInplaceOrView);
```
这对应了"推理模式下不做视图追踪、不记录视图-基础张量关系"的性能收益；离开作用域时则把 ADInplaceOrView 重新加回 included 集合，从而恢复对"推理张量被原地修改"这一非法操作的拦截能力（源码注释第 1 点解释了为什么非推理模式下该 key 必须存在：没有 ADInplaceOrView kernel，就无法对 inference_tensor.add_(1) 抛出错误）。

把 autograd 相关的 DispatchKey 集合整体放入 excluded 集合（启用时）：

DispatchKeySet excluded = enabled
    ? (prev_keyset.excluded_ | c10::autograd_dispatch_keyset)
    : (prev_keyset.excluded_ - c10::autograd_dispatch_keyset);

禁止拷贝/移动。构造函数与析构函数保存并恢复 AutogradState 与 LocalDispatchKeySet，且声明 const/&& 重载为 delete，杜绝守卫被意外复制或移动导致状态错乱（c10/core/InferenceMode.h）。

在 torch::autograd 命名空间中，InferenceMode 被直接别名到 c10::InferenceMode（见 torch/csrc/autograd/InferenceMode.h），因此两种限定写法等价。

推理张量（Inference Tensors）规则

InferenceMode 作用于一段代码块之后，该块内新分配的非视图张量都会被标记为 inference tensors。推理张量遵循两条关键约束：

没有版本计数器：推理张量不携带版本号，因此任何试图读取其版本的行为都会触发错误——典型的触发场景是你把该张量 save_for_backward 保存下来供反向使用；
在 InferenceMode 之外不可变：离开 InferenceMode 作用域后，若尝试——
- 修改（mutate）推理张量的数据；
- 将推理张量置为 requires_grad = true；
都会抛出错误。文档给出的推荐绕行方案是：先在 InferenceMode 之外 clone 出一份普通张量，再对新副本进行修改。

判定规则可以精确表述为：

一个非视图张量是推理张量，当且仅当它是在 InferenceMode 内被分配的；
一个视图张量是推理张量，当且仅当它是某个推理张量的视图。

性能保证

在 InferenceMode 块内，文档明确给出了三点性能与语义保证：

与 NoGradMode 一致：即使输入张量 requires_grad=true，所有算子也都不会记录 grad_fn。该保证同时适用于推理张量与普通张量；
不做视图追踪：对推理张量执行 view 算子时不会进行 view tracking，因此视图与非视图的推理张量在行为上不可区分（这也是从调度层面移除 ADInplaceOrView 的直接结果）；
原地操作不触发版本自增：对推理张量的 inplace 算子保证不会发生 version bump。

基础示例：优化推理前向

{
    c10::InferenceMode guard;
    // Optimized inference without gradient tracking
    auto result = model->forward(input);
}

推理工作负载示例：完整推理流水线

文档给出的一个典型推理流水线示例，将"加载 JIT 模型 + 前处理 + 前向 + 后处理"整体包在 InferenceMode 中：

c10::InferenceMode guard;
model.load_jit(saved_model);
auto inputs = preprocess_tensors(data);
auto out = model.forward(inputs);
auto outputs = postprocess_tensors(out);

值得注意：由于推理张量具备"离开 InferenceMode 后不可变"的特性，若后处理逻辑需要对 out 做 inplace 修改，应把它放到守卫作用域之外，或先对结果做一次 clone。

嵌套 InferenceMode：支持开关状态的任意嵌套

与其它某些守卫不同，InferenceMode 支持以不同 enabled 状态进行嵌套。内层守卫会基于外层状态继续切换，并在退出内层时精确恢复到内层进入前的状态：

{
    c10::InferenceMode guard(true);
    // InferenceMode is on
    {
        c10::InferenceMode guard(false);
        // InferenceMode is off
    }
    // InferenceMode is on
}
// InferenceMode is off

从实现上理解，这一嵌套能力来自两个保存/恢复机制：构造函数记录了进入前的 prev_mode（AutogradState）与 prev_keyset（LocalDispatchKeySet），而 AutogradState 又是一个可整体赋值替换的 POD 式结构（c10/core/AutogradState.h），因此任意层级的开关组合都可以精确入栈出栈。

InferenceMode vs NoGradMode：关键差异

对于纯推理负载，官方文档的结论非常明确：InferenceMode 优先于 NoGradMode，因为它提供更好的性能。两者的关键差异可归纳为：

对比维度	NoGradMode / NoGradGuard	InferenceMode
是否记录 `grad_fn`	关闭	关闭
是否影响张量创建	不影响，创建的仍是普通张量	影响，新分配张量被标记为推理张量
视图追踪 / 版本计数器	照常进行	跳过（性能收益来源）
作用域外对模式内张量的限制	无	推理张量不可变、不可置 `requires_grad=true`、无版本计数器
嵌套开关	受限	支持以 enabled/disabled 状态任意嵌套
适用场景	需要临时关梯度、张量后续还要参与 autograd	确定不会与 autograd 交互的纯推理代码

一句话概括：两者都"跳过与推理无关的执行工作"，但 InferenceMode 还改变了张量的创建方式，NoGradMode 则不会——这正是推理张量限制以及额外性能收益的来源。

从 AutoNonVariableTypeMode 安全迁移

遗留的 AutoNonVariableTypeMode 守卫（现已改名为 AutoDispatchBelowADInplaceOrView）过去常被用于推理负载，但文档明确指出它是不安全的：它会静默绕过安全检查并可能产生错误结果。迁移时应按场景二选一：

场景一：纯推理负载 → 使用 c10::InferenceMode

对于只做推理的工作负载（典型如：加载预训练 JIT 模型并在 C++ 运行时中执行推理），c10::InferenceMode 可作为即插即用（drop-in）的替代品——它保留了原有性能特征，同时提供正确性保证：

c10::InferenceMode guard;
model.load_jit(saved_model);
auto inputs = preprocess_tensors(data);
auto out = model.forward(inputs);
auto outputs = postprocess_tensors(out);

场景二：自定义 autograd kernel → 使用 AutoDispatchBelowADInplaceOrView

对于需要在 Autograd DispatchKey 之下重新分发（redispatch）的自定义 autograd kernel，则应改用 AutoDispatchBelowADInplaceOrView。文档给出的典型示例是一个自定义的 ROIAlign 前向：

class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
 public:
  static torch::autograd::variable_list forward(
      torch::autograd::AutogradContext* ctx,
      const torch::autograd::Variable& input,
      const torch::autograd::Variable& rois,
      double spatial_scale, int64_t pooled_height,
      int64_t pooled_width, int64_t sampling_ratio, bool aligned) {
    ctx->saved_data["spatial_scale"] = spatial_scale;
    ctx->save_for_backward({rois});
    at::AutoDispatchBelowADInplaceOrView guard;
    auto result = roi_align(input, rois, spatial_scale,
        pooled_height, pooled_width, sampling_ratio, aligned);
    return {result};
  }
};

可以看到，这里正确的意图是在自定义 autograd 前向中暂时绕过 Autograd 层的再次分发去调用底层算子，而不是"进入推理模式"——因此必须使用 at::AutoDispatchBelowADInplaceOrView，而不是把它误替换成 c10::InferenceMode（后者会把该 kernel 内的张量标记为推理张量，带来不可变的额外限制）。

总结：如何选择梯度模式守卫

推理代码与训练/求梯度代码交织，只需在个别片段临时关梯度 → torch::NoGradGuard；
整段代码确定与 autograd 零交互（加载模型、纯前向、离线推理），希望获得视图追踪与版本计数方面的额外性能收益，并接受推理张量的不可变约束 → c10::InferenceMode；
自定义 autograd 函数内部需要跳过 Autograd 层分发、直连底层实现 → at::AutoDispatchBelowADInplaceOrView（不要再使用已废弃且不安全的 AutoNonVariableTypeMode 命名）。

若希望进一步深入理解实现细节，可在当前仓库中查看以下源码：

c10/core/InferenceMode.h：守卫的完整实现与设计注释（TLS 状态、DispatchKeySet 切换、不可拷贝约束）；
c10/core/AutogradState.h：被守卫保存/恢复的线程局部 autograd 状态结构；
c10/core/GradMode.h：NoGradGuard 对 AutoGradMode 的继承定义；
aten/src/ATen/core/grad_mode.h 与 torch/csrc/autograd/grad_mode.h：GradMode/AutoGradMode/NoGradGuard 在 at 与 torch::autograd 命名空间中的别名导出；
torch/csrc/autograd/InferenceMode.h：torch::autograd::InferenceMode = c10::InferenceMode 的别名声明。