详解 torch.no_grad() 与 inference_mode() 的底层实现差异及选型建议

在 PyTorch 的模型推理阶段，我们通常需要禁用梯度计算，以节省内存和提高运行速度。最常用的方法是使用上下文管理器 torch.no_grad()。然而，PyTorch 1.9 版本引入了一个更强大的替代品：torch.inference_mode()。本文将详细解析这两者的底层差异，并提供选型建议。

1. 为什么需要禁用梯度？

当我们进行前向传播时，PyTorch 默认会记录操作历史，以便在反向传播时计算梯度（通过 Autograd 机制）。在推理阶段，由于不需要更新权重，这些历史记录是冗余的。禁用 Autograd 可以带来两大好处：

1. 节省内存： 不存储中间激活的梯度信息。
2. 提高速度： 避免了构建和维护计算图的开销。

2. 经典工具：torch.no_grad()

torch.no_grad() 是 PyTorch 中最常用的推理上下文管理器。它的核心功能是修改当前线程的 Autograd 状态，确保在 with 块内创建的任何张量都不会追踪其操作历史 (requires_grad 属性虽然可以为 True，但不会真正记录操作)。

底层实现： 它主要通过 C++ 的 at::ThreadLocalState 来设置一个标志位，告诉 Autograd 系统“不要记录操作”。

限制： 尽管 no_grad() 禁用了梯度计算，但它没有禁用张量的版本计数（version counter）和张量修改历史检查。这意味着如果你在推理过程中进行了大量原地（in-place）操作，系统仍然会花费时间检查这些修改是否可能破坏计算图（尽管我们知道此时不需要计算图）。

3. 现代优化：torch.inference_mode()

torch.inference_mode() 是为纯推理场景量身定制的优化工具，从 PyTorch 1.9 版本开始引入。它提供了比 no_grad() 更深层次的优化。

核心差异： inference_mode() 不仅像 no_grad() 一样禁用了梯度跟踪，它还禁用了：

1. 张量版本计数 (Version Counters)： 用于检测在反向传播时是否发生了不安全的原地修改。
2. In-place 历史跟踪： 彻底停止对张量原地修改的记录和检查。

通过禁用这些额外的检查，inference_mode() 能够进一步减少运行时开销，尤其是在涉及大量张量操作或原地操作（如 ReLU Inplace）时，性能提升更为显著。

4. 实操对比与性能测试

我们通过一个简单的代码示例来比较两者的用法和性能差异。

import torch
import time

# 1. 准备模型和数据
# 使用一个简单的卷积模型进行测试
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 32, 3, padding=1),
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(2)
).cuda() # 假设使用GPU加速

# 将模型设置为推理模式
model.eval()

data = torch.randn(16, 3, 224, 224, device='cuda')

ITERATIONS = 500

def run_benchmark(ctx, name):
    # 预热
    for _ in range(10): 
        with ctx: 
            _ = model(data)

    torch.cuda.synchronize()
    start_time = time.time()

    for _ in range(ITERATIONS):
        with ctx:
            _ = model(data)

    torch.cuda.synchronize()
    end_time = time.time()
    print(f"{name} 平均耗时: {(end_time - start_time) / ITERATIONS * 1000:.4f} ms")

# 2. 运行对比

# 运行 no_grad()
run_benchmark(torch.no_grad(), "torch.no_grad()")

# 运行 inference_mode()
run_benchmark(torch.inference_mode(), "torch.inference_mode()")

# 3. 结果观察（示例输出，实际数据取决于硬件）
# torch.no_grad() 平均耗时: 1.2500 ms
# torch.inference_mode() 平均耗时: 1.1800 ms
# 可以观察到 inference_mode 通常会带来轻微但持续的性能优势。

5. 选型建议

推荐原则：

1. 对于现代 PyTorch (>= 1.9) 的生产推理场景： 始终使用 torch.inference_mode()。它提供了最彻底的性能优化，是官方推荐的推理标准。
2. 对于需要暂时禁用 Autograd 但之后仍然可能需要进行操作追踪的复杂训练流程 (例如，GANs 中只更新 D 而不更新 G 时)： 仍然可以使用 torch.no_grad()，因为它对 Autograd 系统的侵入性更小，更容易恢复状态。

总之，在 PyTorch 推理加速中，torch.inference_mode() 是实现低延迟和低内存占用的首选利器。