PyTorch inference_mode 和 no_grad 怎么选？检查梯度与推理张量

PyTorch 推理时，no_grad() 和 inference_mode() 都能关闭反向图记录。选择的关键是：本次产生的张量，之后是否还要作为可训练计算的输入。只做推理输出可考虑 inference_mode()；特征之后要传给可训练层时，优先用 no_grad() 并验证梯度链路。

本文依据 PyTorch 2.14 官方说明，并在 Windows、Python 3.11、PyTorch 2.14.1 CPU 上运行小张量示例。示例验证梯度行为与张量限制，没有测速度、显存或真实模型效果，不能据此断言其中一种方式总是更快。

先把三个开关分清

no_grad 接口说明提醒，显式接收 requires_grad 的工厂函数有例外，所以“进入这个区域后所有张量必然都不需要梯度”也不准确。本文用的是普通前向计算，不涉及这个例外。

inference_mode 接口说明明确指出它不会替你调用 model.eval()。正式推理时通常需要分别完成评估模式设置和梯度模式选择。

可运行对比：同一特征接入可训练层

在独立 Python 环境安装 CPU 依赖，把完整代码保存为 inference_mode_demo.py，执行 python inference_mode_demo.py。这里的线性层随机初始化，只用于演示接口。

python -m pip install torch==2.14.1 numpy --index-url https://download.pytorch.org/whl/cpu

import torch
from torch import nn

torch.manual_seed(7)
features = nn.Linear(2, 3).eval()
x = torch.tensor([[1.0, 2.0]])

with torch.no_grad():
    ordinary = features(x)
with torch.inference_mode():
    inference = features(x)

torch.testing.assert_close(ordinary, inference)
print('requires_grad:', ordinary.requires_grad, inference.requires_grad)
print('is_inference:', torch.is_inference(ordinary), torch.is_inference(inference))

# 取出的普通张量可以用作后续可训练层的输入。
head = nn.Linear(3, 1)
head(ordinary).sum().backward()
assert head.weight.grad is not None
print('ordinary_to_trainable_head:', True)

# 推理张量在这里需要被保存以计算 head 权重梯度，触发限制。
head.zero_grad(set_to_none=True)
blocked = False
try:
    head(inference).sum().backward()
except RuntimeError as exc:
    if 'Inference tensors cannot be saved for backward' not in str(exc):
        raise
    blocked = True
assert blocked
print('inference_to_trainable_head_blocked:', blocked)

# 在 inference_mode 区域外复制，得到可供自动求导使用的普通张量。
ordinary_copy = inference.clone()
assert not torch.is_inference(ordinary_copy)
head.zero_grad(set_to_none=True)
head(ordinary_copy).sum().backward()
assert head.weight.grad is not None
print('outside_clone_to_trainable_head:', True)

# 关闭梯度不会替你设置模型的评估模式。
dropout = nn.Dropout(p=0.5)
with torch.inference_mode():
    assert dropout.training
dropout.eval()
with torch.inference_mode():
    result = dropout(torch.ones(8))
assert torch.equal(result, torch.ones(8))
print('eval_dropout_is_identity:', True)

前半段用两种模式计算同一份特征并比较输出，后半段用一个新的可训练线性层做反向计算。这个线性层为了求权重梯度，需要保存输入；因此直接使用推理张量会触发限制。自动求导机制说明给出的选择条件也是：是否打算把推理模式中创建的张量放进之后需要自动求导的计算。

看标记和梯度，不靠猜测

上述 CPU 示例实际得到：

requires_grad: False False
is_inference: False True
ordinary_to_trainable_head: True
inference_to_trainable_head_blocked: True
outside_clone_to_trainable_head: True
eval_dropout_is_identity: True

两份特征都不需要梯度，却只有后一份属于推理张量；因此仅检查 requires_grad=False 不足以发现边界。普通特征作为输入时，可训练层的权重有梯度；推理特征在这个需要保存输入的计算中被拒绝。把它在推理区域外 clone() 成普通张量后，这条反向计算才通过。

复制会增加内存和数据搬运开销，也不能恢复之前被关闭的梯度链。如果原本希望连特征提取层一起训练，就应该让它的前向处于正常梯度模式；不要靠事后复制来补回不存在的计算图。

接回自己的推理或训练流程

1. 仅输出分数、文本、分类结果：先 model.eval()，在 inference_mode() 内前向；核对输出形状、有限值和任务结果。
2. 固定特征提取器，但后面接一个需要训练的层：用 no_grad() 提取特征，随后退出该区域再训练新层，检查新层参数的 grad 是否存在。
3. 需要输入梯度、解释方法或训练原模型：前向必须保留所需自动求导，不使用这两种关闭梯度模式覆盖该计算。

遇到 Inference tensors cannot be saved for backward，沿数据流查找推理区域中新建的张量。若后续确实要训练，优先调整到 no_grad()；必须复制时在区域外复制，并验证梯度只覆盖预期部分。不要用捕获所有异常后继续运行来掩盖训练失败。

如需比较速度，在同一模型、相同输入、相同设备和评估模式下单独测量，处理预热、同步与重复采样。这里的行为演示只回答该选哪种梯度模式，没有提供性能结论。