PyTorch 推理工程（03）：train、eval、no_grad 与 inference_mode

1. 本节定位

常见推理脚本同时包含 model.eval() 与 torch.inference_mode()（或 no_grad()）：

1
2
3
4
model.eval()

with torch.inference_mode():
    y = model(x)

二者作用于不同子系统：eval() 调整部分模块（如 BatchNorm、Dropout）的行为；inference_mode() / no_grad() 控制是否构建 autograd 图。仅设置其一通常不足以同时满足数值行为与开销两方面的预期，混用不当也会出现与视图、原地操作相关的报错。

2. 模块模式与 autograd 模式（分类图）

在展开讲之前，先把这 4 个东西的关系说清楚。

它们其实分属两个完全独立的系统：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
┌─────────────────────────────────────────────────────┐
│              两个独立的控制维度                        │
│                                                     │
│  维度 1：模块行为模式（某些层怎么跑）                   │
│  ┌──────────────────────────────────────────────┐   │
│  │  model.train()   ←→   model.eval()           │   │
│  │  影响：Dropout, BatchNorm 等层的前向行为        │   │
│  └──────────────────────────────────────────────┘   │
│                                                     │
│  维度 2：autograd 模式（PyTorch 要不要记录计算过程）   │
│  ┌──────────────────────────────────────────────┐   │
│  │  默认 grad mode                              │   │
│  │  torch.no_grad()                             │   │
│  │  torch.inference_mode()                      │   │
│  └──────────────────────────────────────────────┘   │
│                                                     │
└─────────────────────────────────────────────────────┘

这两个维度互不干涉：

  • eval() 不会关闭梯度记录
  • no_grad() 不会让 Dropout 停下来

两维相互独立：仅 eval() 不能关闭 autograd；仅 no_grad() 不能改变 Dropout/BatchNorm 等行为。典型推理路径需同时设置模块模式与 autograd 模式。

后文均围绕下列关系展开：

train() / eval() 控制的是模块行为模式；no_grad() / inference_mode() 控制的是 autograd 记录行为。两者是独立的，不能互相替代。

3. 背景：PyTorch 训练时在做什么额外的事

训练前向与计算图

训练分为前向与反向两阶段：前向计算 loss；反向沿计算图求梯度并更新参数。

在默认梯度模式下，前向过程会同步构建供 .backward() 使用的计算图：

具体来说，它会在内存里构建一张计算图（computation graph），记录：

  • 每一步做了什么操作
  • 每个中间结果从哪里来
  • 这些信息是为了后面 .backward() 时能顺着图求导
1
2
3
前向：x → Linear → ReLU → Linear → loss
                 ↓ 同时构建计算图 ↓
反向：loss ← 求导 ← ReLU ← Linear ← x（利用计算图）

推理与 autograd

推理仅需 x → 模型 → y，不求梯度、不更新权重。若仍启用默认梯度模式，将承担构图带来的内存与开销；故推理路径中通常使用 torch.no_grad() 或 torch.inference_mode() 关闭记录。

1
2
with torch.no_grad():
    y = model(x)

4. train() 和 eval()：控制模块行为

train() / eval() 切换模块的训练/评估状态；对 Linear、ReLU、LayerNorm 等多数算子无行为差异。

只有少数层在两种模式下行为不同，主要是：

  • Dropout
  • BatchNorm

下面分别说说。

5. Dropout：train/eval 下行为差异

Dropout 的原理是：训练时随机"关掉"一部分神经元（把它们的输出置零），强迫模型不依赖特定的某几个神经元，起到正则化的效果。

但推理时不需要这个随机性：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import torch
import torch.nn as nn

drop = nn.Dropout(p=0.5)    # 每个元素有 50% 概率被置零
x = torch.ones(10)          # 全 1 输入

# 训练模式：随机置零
drop.train()
print(drop(x))   # 例：tensor([2., 0., 2., 0., 2., 0., 0., 2., 0., 2.])
print(drop(x))   # 例：tensor([0., 2., 0., 2., 0., 2., 2., 0., 2., 0.])
# 每次不一样！（而且没被置零的元素会乘以 1/(1-p)=2 补偿）

# 评估模式：原样通过
drop.eval()
print(drop(x))   # tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
print(drop(x))   # tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
# 每次都一样，不丢任何值

注意：训练模式下非零位置的值是 2.0，不是 1.0。这是 Dropout 的"缩放补偿"：为了让期望值一致，存活的元素会除以 (1-p)，这里 1/(1-0.5)=2。

结论：含 Dropout 的模型在推理前应调用 model.eval()，否则同一输入可能因随机掩码而得到不同输出。

6. BatchNorm：train/eval 下行为差异

BatchNorm 相对复杂，但推理岗位必须理解。

它做了什么

BatchNorm 对每个 batch 的输入做归一化：把均值调成 0，方差调成 1，然后用可学习的参数 γ（gamma，weight）和 β（beta，bias）还原尺度。

训练时和推理时的区别

阶段 用什么统计量归一化 会不会更新 running 统计量
训练 当前 batch 自己的均值/方差 会，每个 batch 都更新
推理 训练时积累的 running_mean / running_var 不会 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torch
import torch.nn as nn

bn = nn.BatchNorm1d(4)
x = torch.randn(8, 4)   # batch of 8 samples，每个 4 维

# 训练时：用当前 batch 的统计量
bn.train()
out = bn(x)
print(bn.running_mean)  # 每次 forward 后都会更新

# 推理时：用积累下来的统计量
bn.eval()
out = bn(x)
print(bn.running_mean)  # 不再更新，保持训练时积累的值

结论：如果推理时忘了 model.eval()，BatchNorm 会用当前 batch（可能只有 1 条数据）的统计量做归一化，而不是训练时学到的统计量，结果往往是错的。

这是"模型能跑但推理结果不对"的经典原因之一。

7. eval() 只影响模块行为，不影响梯度

这是最容易误解的点，直接用代码证明：

1
2
3
4
5
6
7
8
9
10
import torch
import torch.nn as nn

model = nn.Linear(4, 2)   # Linear 层默认参数 requires_grad=True
x = torch.randn(3, 4)

model.eval()              # 切换到评估模式
y = model(x)              # 前向推理

print(y.requires_grad)    # True ← ！梯度追踪依然开着

即使 model.eval() 了，输出依然有 requires_grad=True，说明 autograd 依然在工作。

因此仅调用 eval() 不足以关闭 autograd，仍需显式进入 no_grad 或 inference_mode：

1
2
3
4
5
model.eval()
with torch.no_grad():
    y = model(x)

print(y.requires_grad)    # False ← 现在才关掉了

8. torch.no_grad()：关闭 autograd 记录

torch.no_grad() 是一个上下文管理器，进入这个块后，PyTorch 不再构建计算图：

1
2
3
4
5
6
7
8
9
import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# x 本来要参与梯度计算

with torch.no_grad():
    y = x * 2        # 这个计算不会被记录进计算图

print(y.requires_grad)   # False ← 即使 x 有 grad，y 也没有

在推理里的实际意义

模型参数默认 requires_grad=True。未进入 no_grad() / inference_mode() 时：

1
2
model.eval()
y = model(x)   # 所有中间激活值都被加入了计算图，占用额外内存

进入 no_grad() 后：

1
2
3
model.eval()
with torch.no_grad():
    y = model(x)   # 所有中间激活值不记入计算图，内存大幅减少

对于一个有很多层的大模型，no_grad() 带来的内存节省是相当可观的——因为不需要保存反向传播所需的中间激活值了。

9. torch.inference_mode()：更适合纯推理的选择

inference_mode() 是 PyTorch 1.9 之后引入的，可以理解为"推理场景专用的 no_grad"。

1
2
with torch.inference_mode():
    y = model(x)

它和 no_grad() 的共同点

  • 都不构建 backward 计算图
  • 都适合推理
  • 都能减少计算和内存开销

它比 no_grad() 多做的事

PyTorch 内部维护着很多状态，帮助 autograd 工作，比如：

  • version counter：追踪一个 Tensor 被 in-place 修改了几次（为了保证反向传播的正确性）
  • grad_fn 记录：记录每个 Tensor 的计算来源

no_grad() 只是不构建计算图，但上面这些内部跟踪机制还在运行。

inference_mode() 则更激进：直接告诉 PyTorch，这里创建的 Tensor 是纯推理产物，所有 autograd 相关的跟踪机制都不需要。

所以 inference_mode() 通常比 no_grad() 省更多开销。

代价：更严格的限制

在 inference_mode() 里创建的 Tensor，离开这个上下文后，不能再参与需要 autograd 的计算：

1
2
3
4
5
6
with torch.inference_mode():
    y = model(x)

# y 是 inference tensor，不能放回 autograd 计算图
loss = y.sum()
loss.backward()   # 反例: 可能报错：y 是 inference tensor

而用 no_grad() 产生的 Tensor 则没有这个限制（虽然它本身没有 grad，但可以参与后续有 grad 的计算）：

1
2
3
4
with torch.no_grad():
    y = model(x)

# 可以继续用，例如 KL divergence 计算（虽然不常见）

10. 一张表：四种组合的效果对比

场景 Dropout 行为 BatchNorm 行为 梯度追踪 额外 autograd 开销
model.train() 默认 随机丢弃 用 batch 统计 是 是
model.eval() 默认 全部保留 用 running 统计 是 是
model.eval() + no_grad() 全部保留 用 running 统计 否 部分节省
model.eval() + inference_mode() 全部保留 用 running 统计 否 更省

推理时想要的是最后一行：Dropout 不丢弃、BatchNorm 用 running 统计、不追踪梯度、最小开销。

11. 标准推理代码的两个模板

模板 1：通用版（兼容性最好）

1
2
3
4
model.eval()

with torch.no_grad():
    y = model(x)

适合：

  • 大多数推理代码
  • 推理结果后面还可能参与其他计算
  • 兼容性要求高、不想踩 inference_mode 的兼容坑

模板 2：纯推理版（性能最优）

1
2
3
4
model.eval()

with torch.inference_mode():
    y = model(x)

适合：

  • 纯前向推理，不需要后续梯度操作
  • 追求更低延迟和更少内存
  • 工程部署中的标准推理服务

12. 完整对比实验：自己跑一遍

建议把下面这个脚本跑一遍，把每个输出和原理对应起来：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(8, 2)
)

x = torch.randn(3, 4)

# ── 场景 1：train 模式，默认 grad ──
model.train()
y1 = model(x)
y2 = model(x)
print("【场景1: train模式】")
print("  两次输出相同？", torch.allclose(y1, y2))   # 大概率 False（Dropout 随机）
print("  requires_grad?", y1.requires_grad)          # True

# ── 场景 2：eval 模式，默认 grad ──
model.eval()
y3 = model(x)
y4 = model(x)
print("\n【场景2: eval模式，无no_grad】")
print("  两次输出相同？", torch.allclose(y3, y4))   # True（Dropout 关了）
print("  requires_grad?", y3.requires_grad)          # True  ← 注意！梯度还在

# ── 场景 3：eval + no_grad ──
model.eval()
with torch.no_grad():
    y5 = model(x)
print("\n【场景3: eval + no_grad】")
print("  requires_grad?", y5.requires_grad)          # False ← 关了

# ── 场景 4：eval + inference_mode ──
model.eval()
with torch.inference_mode():
    y6 = model(x)
print("\n【场景4: eval + inference_mode】")
print("  requires_grad?", y6.requires_grad)          # False
print("  is_inference?", y6.is_inference())          # True ← 被标记为推理 tensor

预期输出（大致）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
【场景1: train模式】
  两次输出相同？ False
  requires_grad? True

【场景2: eval模式，无no_grad】
  两次输出相同？ True
  requires_grad? True

【场景3: eval + no_grad】
  requires_grad? False

【场景4: eval + inference_mode】
  requires_grad? False
  is_inference? True

y6.is_inference() 为 True，说明这个 Tensor 被标记为"推理产物"，PyTorch 知道不能把它拿回 autograd 计算图里用。

13. requires_grad 和模型参数的关系说明

初学者常见的困惑：我没有手动设 requires_grad，为什么输出有梯度跟踪？

原因：nn.Linear 等层的参数（weight、bias）默认 requires_grad=True。

若参与运算的参数 requires_grad=True，即使输入未要求梯度，输出仍可能接入 autograd 计算链。

1
2
3
4
5
6
linear = nn.Linear(4, 2)
print(linear.weight.requires_grad)   # True（参数默认可学习）

x = torch.randn(3, 4)               # x 本身没有 requires_grad
y = linear(x)
print(y.requires_grad)               # True！因为链路上有 requires_grad=True 的参数

因此即使输入 x 的 requires_grad 为假，只要参数需追踪，输出仍可携带梯度信息。

而加了 no_grad() / inference_mode() 后，这条追踪链就被整个切断了。

14. 推理脚本检查清单

审阅推理脚本时可依次核对：

  • [ ] 推理前有没有 model.eval()？
    → 没有的话，Dropout 会随机丢值，BatchNorm 用 batch 统计

  • [ ] 前向外面有没有 torch.no_grad() 或 torch.inference_mode()？
    → 没有的话，有不必要的计算图构建开销和内存占用

  • [ ] 两件事有没有都做，还是只做了其中一个？
    → 只写 eval() 不够；只写 no_grad() 也不够

  • [ ] 如果用了 inference_mode()，推理结果有没有拿去参与后续梯度计算？
    → 有的话要换成 no_grad()

15. 常见误区

误区 1：eval() 等于关闭梯度

错。eval() 只切换模块的行为模式，不碰 autograd。

误区 2：no_grad() 会让 Dropout 停下来

错。no_grad() 只关 autograd，Dropout 的 train/eval 行为由模块模式决定，和 autograd 无关。

误区 3：推理只写 eval() 就完整了

不完整。还需要 no_grad() 或 inference_mode() 来关闭不必要的梯度追踪。

误区 4：inference_mode() 就是 no_grad() 的别名

不是。inference_mode() 更激进，额外关闭了更多 autograd 跟踪机制，产生的 Tensor 不能再参与 autograd 计算。

误区 5：模型不在训练，所以 grad 自然不会记录

否。默认梯度模式下前向会被 autograd 追踪，与是否处于训练业务逻辑无关。

16. 与推理模式相关的辨析

Q：model.eval() 做了什么？
A：把模块切换到评估模式，影响的是那些训练和推理行为不同的层（主要是 Dropout 和 BatchNorm）。不影响梯度记录。

Q：为什么 eval() 不能代替 torch.no_grad()？
A：两者控制的是完全不同的系统：eval() 控制模块行为，no_grad() 控制 autograd 记录。缺一不可。

Q：torch.no_grad() 的作用是什么？
A：让这一段计算不被记录到 backward 计算图，忽略不必要的梯度追踪开销，减少内存占用。

Q：torch.inference_mode() 和 torch.no_grad() 的区别是什么？
A：都不构建计算图；但 inference_mode() 更激进，还关闭了更多内部 autograd 跟踪机制（如 version counter），性能开销更小，但其中创建的 Tensor 不能再参与 autograd 计算。

Q：标准推理代码为什么要同时写 model.eval() 和 torch.inference_mode()？
A：前者让 Dropout 不丢弃、BatchNorm 用 running 统计；后者让推理过程不做任何梯度相关的额外工作。二者分属不同系统，必须都写。

17. 思考题

建议自己动手跑，先猜输出，再验证。

练习 1：验证 eval() 不会关闭梯度

1
2
3
4
5
6
7
model = nn.Linear(4, 2)
x = torch.randn(3, 4)

model.eval()
y = model(x)

print(y.requires_grad)   # 猜一下：True 还是 False？

思考：为什么即使没有人设置 requires_grad=True，输出还是有梯度跟踪？

练习 2：验证 Dropout 受 train/eval 控制，不受 no_grad 控制

1
2
3
4
5
6
7
8
9
10
11
drop = nn.Dropout(p=0.5)
x = torch.ones(10)

# 情况 A：train 模式 + no_grad
drop.train()
with torch.no_grad():
    print(drop(x))   # 猜：有没有被置零？

# 情况 B：eval 模式，无 no_grad
drop.eval()
print(drop(x))       # 猜：有没有被置零？

思考：哪种情况下 Dropout 起作用？这说明了什么？

练习 3：验证 no_grad() 切断梯度链

1
2
3
4
5
6
7
8
9
10
11
linear = nn.Linear(4, 2)
x = torch.randn(3, 4)

# 情况 A
y1 = linear(x)
print(y1.requires_grad)   # ?

# 情况 B
with torch.no_grad():
    y2 = linear(x)
print(y2.requires_grad)   # ?

思考：linear.weight 的 requires_grad 变化了吗？为什么 y2 没有 grad？

练习 4：inference_mode 的 is_inference 标记

1
2
3
4
5
6
7
8
9
10
11
model = nn.Linear(4, 2)
model.eval()

with torch.inference_mode():
    y = model(x)

print(y.is_inference())    # True 还是 False？

# 离开 inference_mode 后，尝试对 y 做一个有梯度的操作
z = y.sum()
# z.backward()   # 取消注释，观察是否报错

思考：inference_mode 比 no_grad 多了哪些限制？为什么这些限制让它更快？

练习 5：写一个标准推理函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def infer(model, x, device="cpu"):
    """
    标准推理函数模板：
    1. 把输入移到正确设备
    2. 切换模型行为模式
    3. 关闭梯度追踪
    4. 执行前向
    5. 返回结果
    """
    # （补全）

model = nn.Linear(4, 2)
x = torch.randn(3, 4)
y = infer(model, x)
print(y.shape, y.requires_grad)

进阶：让这个函数同时接受 device 参数，自动把 model 和 x 移到指定设备。

18. 本节要点与自检

  • 能清晰区分"模块行为模式"和"autograd 模式"，知道它们是两个独立维度
  • 知道 Dropout 和 BatchNorm 为什么在 train/eval 下行为不同
  • 知道 eval() 为什么不能代替 no_grad()
  • 知道 no_grad() 为什么不能代替 eval()
  • 理解 inference_mode() 比 no_grad() 多做了什么，以及对应的限制
  • 能一眼看出一个推理脚本是否缺了关键步骤
  • 能用标准模板正确写出纯推理函数

19. 小结

eval() 负责让模型"像推理一样工作"，no_grad() / inference_mode() 负责让 PyTorch"像推理一样执行"。两件事独立，推理时都要做。

系列导航