第39篇|autograd 自动求导
本篇定位:专栏四「PyTorch 与 AI Infra 入门」的第三篇。承接第 38 篇的 GPU、
device与 CUDA,把“算子在哪里执行”继续推进到“梯度如何沿计算路径返回”。读完你将能够:用
requires_grad控制求导范围;理解计算图、backward()、.grad、detach()、no_grad()与inference_mode()的关系;处理梯度累积、非标量输出、重复反向和显存释放;把第 36 篇 NumPy 手写反向传播的公式映射到 PyTorch,并从 AI Infra 角度分析训练与推理的开销。预计用时:阅读 50~70 分钟,动手 60~90 分钟。
第 38 篇里,我们把 Tensor 放到了正确的 device:数据在哪里,算子就在哪里执行。但训练一个模型不只需要一次前向计算。模型先用参数产生预测,再用损失衡量误差,最后还要知道每个参数应该向哪个方向移动。这个“从结果沿依赖关系向输入返回影响程度”的过程,就是自动求导(automatic differentiation,简称 autograd)。
第 36 篇用 NumPy 手写 MLP 时,我们显式保存中间值,再用链式法则计算 dW、db。那种写法非常适合理解数学;真实模型却可能有数百亿个参数、分支连接、残差路径和 CUDA kernel。PyTorch 的 autograd 会在前向时记录必要的操作,在反向时按图执行局部导数。它不是魔法,也不是“重新运行一遍 Python”;它是一套围绕 Tensor、计算图和设备内存的运行时机制。
本文的示例优先保持真实 PyTorch 语义。若当前环境未安装 torch,仍可阅读、做语法检查;安装 PyTorch 后,CPU 示例即可运行,CUDA 示例请沿用第 38 篇的可用性检查。
一、先建立计算图心智模型
1.1 从链式法则到依赖图
设 x 是输入,经过两个操作得到 y:
a = x * 2
b = a + 1
y = b ** 2
数学上,y 依赖 b,b 依赖 a,a 依赖 x。如果要求 dy/dx,就沿相反方向相乘局部导数:dy/db = 2b,db/da = 1,da/dx = 2,因此 dy/dx = 4b。计算图(computational graph)就是把这些依赖关系表示成节点和边:Tensor 是数据,操作是产生新 Tensor 的节点,边携带前向产生的值和反向所需的上下文。
PyTorch 默认使用动态计算图。一次前向执行到哪里,图就构建到哪里;下一次前向可以走完全不同的 Python 分支。这和“先把整个程序编译成固定图”不同,也很适合调试与写有条件分支的模型。代价是训练时需要保存中间结果,并在反向后回收图。
import torch
x = torch.tensor(2.0, requires_grad=True)
a = x * 2
b = a + 1
y = b ** 2
print(x.requires_grad) # True
print(a.requires_grad) # True
print(y.grad_fn) # 非 None,表示 y 有反向节点
y.backward()
print(x.grad) # tensor(20.),因为 b=5,dy/dx=2*b*2=20
grad_fn 是观察图的一个入口。对用户创建且需要梯度的叶子 Tensor(leaf Tensor),它通常没有 grad_fn;对运算结果,grad_fn 会指向对应的反向函数,例如乘法、加法或幂运算。不要把 grad_fn 当成完整的可视化图,它只是当前节点对 autograd 引擎的接口。
1.2 requires_grad 是求导边界
只有被纳入求导追踪的路径,才会为反向准备上下文。最常见的入口是创建 Tensor 时指定 requires_grad=True,或者对已有 Tensor 调用 requires_grad_():
import torch
x = torch.tensor([1.0, 2.0, 3.0])
print(x.requires_grad) # False
x.requires_grad_()
y = (x * x).sum()
y.backward()
print(x.grad) # tensor([2., 4., 6.])
在训练中,模型参数通常已经是 requires_grad=True 的 nn.Parameter;输入是否求梯度则取决于任务。普通监督学习只需参数梯度,输入可以不求梯度。对输入求梯度常见于对抗样本、可解释性、扩散模型的某些步骤和梯度引导采样。无目的地给所有中间 Tensor 打开梯度,会扩大图和显存占用。
requires_grad 不等于“这个 Tensor 一定有 .grad”。它只是表示从这里开始允许构图。默认情况下,反向结束后,叶子 Tensor 的梯度写入 .grad;非叶子中间结果通常不保留 .grad,除非显式调用 retain_grad()。这一区分能解释很多“我明明看到它 requires_grad,为什么 grad 是 None”的问题。
1.3 叶子 Tensor、结果 Tensor 与 .grad
叶子 Tensor(leaf Tensor)可以粗略理解为不是由当前 autograd 图中某个操作产生、而是由用户或参数存储的 Tensor。参数和显式设置 requires_grad=True 的输入通常是叶子。下面比较三种对象:
import torch
x = torch.tensor(3.0, requires_grad=True) # 叶子
z = x * 2 # 非叶子
z.retain_grad() # 请求保留中间梯度
loss = z ** 2
loss.backward()
print(x.is_leaf, x.grad) # True tensor(24.)
print(z.is_leaf, z.grad) # False tensor(12.)
这里 loss = z²,所以 dloss/dz = 2z = 12,而 dz/dx = 2,最终 dloss/dx = 24。retain_grad() 只影响非叶子 Tensor 的梯度保存,不会改变数学结果,也不会让它成为参数。
AI Infra 视角:
.grad是训练状态,不是普通输出。分布式训练、梯度累积和混合精度都会改变它何时产生、存在哪里、何时同步。排查训练吞吐时,要把“前向激活”“反向梯度”“参数更新”视为三种不同的内存和通信对象。

图1 从前向到反向的计算图
二、backward() 与梯度的真实语义
2.1 backward() 做了什么
对标量 Tensor 调用 backward(),等价于从这个标量开始做反向传播,把沿图可达的梯度累加到叶子节点的 .grad。它不是“返回梯度的函数”那么简单:它会读取保存的中间值,执行反向节点,并把结果写入目标 Tensor 的梯度缓冲区。
import torch
w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
prediction = w * 3 + b
loss = (prediction - 10) ** 2
loss.backward()
print(w.grad) # tensor(-18.)
print(b.grad) # tensor(-6.)
prediction = 7,误差是 -3,损失对 prediction 的导数是 -6。再乘以 dprediction/dw = 3,得到 -18?注意这里 loss = (prediction - 10) ** 2,prediction - 10 = -3,导数是 2 * -3 = -6,所以 w.grad = -18,b.grad = -6。代码注释应与计算一致,实际输出为:
tensor(-18.)
tensor(-6.)
这个例子也提醒我们:不要凭记忆写输出,先根据图验证局部导数。反向传播的每条边都遵循链式法则,复杂模型只是边更多。
2.2 非标量输出需要梯度种子
只有标量可以直接调用无参数的 backward()。如果输出是向量或矩阵,autograd 不知道你想对哪个标量目标求导,需要传入同形状的 gradient,它也叫外部梯度或向量-雅可比乘积(vector-Jacobian product,VJP)中的向量:
import torch
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.backward(torch.ones_like(y))
print(x.grad) # tensor([2., 4.])
这相当于对 y.sum() 求导。若传入 torch.tensor([10.0, 1.0]),结果就是对加权和 10*y[0] + y[1] 求导。训练时损失通常通过 mean() 或 sum() 变成标量,因此最常见的 loss.backward() 不需要显式梯度种子。
不要把 backward() 理解为生成完整雅可比矩阵。大模型训练通常只需要损失对参数的梯度,也就是一次反向的 VJP;完整雅可比会带来巨大的时间和内存成本。需要更高级导数时,应使用 torch.autograd.grad、函数变换或专门的高阶微分 API,并明确成本。
2.3 梯度会累积:清零是训练循环的一部分
PyTorch 默认把新梯度加到已有 .grad,而不是覆盖它。这种设计支持梯度累积(gradient accumulation),但也意味着每个优化步骤前必须清零,除非你明确就是要跨多个 micro-batch 累积:
import torch
w = torch.tensor(1.0, requires_grad=True)
for step in range(2):
loss = (w * 2 - 4) ** 2
loss.backward()
print(step, w.grad.item())
w.grad = None
第一个损失在 w=1 时梯度为 -8。如果不清零,第二次 backward 会把新的 -8 加到旧值上,结果变成 -16。生产训练一般写成 optimizer.zero_grad(set_to_none=True),让梯度缓冲区回到 None,减少不必要的填零写入;若代码依赖一个始终存在的零 Tensor,再选择默认的清零方式。
梯度累积的正确结构是:多个 micro-batch 各自前向和反向,累计到同一组参数;达到目标 batch size 后再执行一次 optimizer.step(),随后清零。若损失使用 mean,还要按 micro-batch 数或样本数正确缩放,否则有效梯度的权重会错。
2.4 和第 36 篇手写反向传播对照
第 36 篇的 NumPy MLP 把 z1、a1、z2 存入缓存,然后手写:
# NumPy 思路:显式写链式法则
# dL/dz2 -> dL/dW2, dL/da1 -> dL/dz1 -> dL/dW1
grad_w2 = a1.T @ grad_z2
grad_a1 = grad_z2 @ w2.T
grad_z1 = grad_a1 * (z1 > 0)
grad_w1 = x.T @ grad_z1
PyTorch 的等价代码只需建立参数和前向表达式,再调用 loss.backward():
import torch
x = torch.tensor([[1.0, 2.0]])
w1 = torch.tensor([[0.2, -0.1], [0.3, 0.4]], requires_grad=True)
w2 = torch.tensor([[0.5], [-0.2]], requires_grad=True)
hidden = torch.relu(x @ w1)
out = hidden @ w2
loss = (out - 1.0).pow(2).mean()
loss.backward()
print(w1.grad)
print(w2.grad)
数学没有消失:ReluBackward、矩阵乘法反向节点和幂运算反向节点正在替你执行局部导数与矩阵乘法。手写版本让你看见公式,autograd 版本让复杂模型可组合。二者应当互相校验:先用小 Tensor、固定值和有限差分验证一个模块,再把它放进 GPU、混合精度或分布式训练。

backward()与梯度累积的关系。前向传播得到prediction和loss,通过loss.backward()沿图反向,生成param.grad的新梯度累积,再调用step()更新参数。下方“梯度账本”说明:第1次梯度为-8,清零后第2次梯度仍为-8;若不清零则累积,若刻意累积则按逻辑batch正确缩放。此图与上文对比手写反向传播,直观呈现autograd自动求导中梯度累积的流程与机制。图2 backward 与梯度累积
三、停止追踪:detach、no_grad 与 inference_mode
3.1 detach() 是 Tensor 级断图
detach() 返回一个不再连接当前计算图的 Tensor。它通常与原 Tensor 共享存储,所以修改其中一个可能影响另一个;它不是深拷贝,也不是把数据搬到 CPU。需要 NumPy 时,CUDA Tensor 还要先 .cpu():
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x * 3
z = y.detach()
print(z.requires_grad, z.grad_fn) # False None
loss = z ** 2
print(loss.requires_grad) # False
detach() 适用于明确的模型边界,例如把教师模型输出作为学生模型的固定目标,把强化学习中的回报从策略图切开,或在生成伪标签时防止梯度回到上游。写出 target = target.detach(),比依赖某个函数“恰好不追踪”更清楚。
如果需要一个完全独立的副本,使用 x.detach().clone();如果还要让副本可求导,再调用 .requires_grad_()。注意顺序和意图:clone() 复制数据,但单独调用 clone() 仍会保留梯度关系;detach().clone() 才是“复制并断开”。
3.2 torch.no_grad() 是代码块级开关
no_grad() 在上下文中关闭梯度记录,适合验证、评估和不需要反向的临时计算:
import torch
w = torch.tensor(2.0, requires_grad=True)
with torch.no_grad():
y = w * 5
print(y.requires_grad) # False
它不会改变参数本身的 requires_grad,离开代码块后,正常前向仍可构图。它也不等于 model.eval():eval() 改变 Dropout、BatchNorm 等模块的行为;no_grad() 改变 autograd 记录。评估模型通常需要两者同时使用:
model.eval()
with torch.no_grad():
scores = model(batch)
反过来,训练时不能用 no_grad() 包住前向,否则 loss 没有反向图。冻结部分参数有多种实现,最简单的是把它们的 requires_grad 设为 False;若只是暂时不让某个分支回传,可以对分支输出 detach()。
3.3 inference_mode() 更强,但边界更严格
torch.inference_mode() 面向纯推理路径,除了关闭梯度记录,还会使用更少的 autograd 元数据,通常能降低开销。它适合在线推理、离线批量打分、验证中完全不需要后续求导的路径:
import torch
with torch.inference_mode():
features = model(batch)
result = features.softmax(dim=-1)
inference_mode() 不是“更快的 eval()”,也不能替代 eval();二者职责仍然不同。它比 no_grad() 更严格,推理模式创建的 Tensor 不适合之后再无缝参与需要 autograd 的计算。如果一个结果未来可能被拿去求导,使用 no_grad() 或显式 detach(),并把边界写在代码中。
在服务系统里,建议把训练、校验、推理三个入口分开:训练入口开启梯度并控制生命周期;校验入口 model.eval() 加 no_grad();纯推理入口 model.eval() 加 inference_mode()。这能让显存监控、延迟基准和故障定位更稳定。

图3 训练、评估与推理的梯度边界
四、图的生命周期、高阶梯度与显存
4.1 为什么第二次 backward() 会报错
默认情况下,反向传播使用完保存的中间结果后会释放它们,以降低峰值内存。对同一张图再次调用 backward(),常见报错是“Trying to backward through the graph a second time”。例如:
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x * x
loss = y + 1
loss.backward()
# loss.backward() # 默认不允许:中间缓存已经释放
如果确实要对同一张图反向多次,可以第一次调用 loss.backward(retain_graph=True)。它保留反向所需的缓存,第二次才有机会继续使用:
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x * x
loss = y + 1
loss.backward(retain_graph=True)
loss.backward()
print(x.grad) # 两次结果累加为 8
这里单次导数是 2x=4,两次累计成 8,不是 16;正确输出应为 tensor(8.)。retain_graph=True 会增加内存,不应成为“报错就加上”的默认修复。更好的办法通常是把需要的损失合并后只反向一次,或在循环中每次重新前向构图。
4.2 创建高阶梯度:create_graph=True
retain_graph 解决“重复使用同一张一阶图”;create_graph=True 解决“让一阶梯度本身也进入新图”。二者目的不同:
import torch
x = torch.tensor(3.0, requires_grad=True)
y = x ** 3
grad_x = torch.autograd.grad(y, x, create_graph=True)[0]
second = torch.autograd.grad(grad_x, x)[0]
print(grad_x, second) # tensor(27.) tensor(18.)
dy/dx = 3x² = 27,二阶导数是 6x = 18。元学习、物理信息网络和某些优化算法需要高阶梯度,但它们会保存更多图和中间量。生产代码应明确记录它的成本,并在 profiler 中观察显存,而不是只看最终 loss 是否正确。
4.3 显存、图引用与循环
只要 Python 变量仍引用计算图中的 Tensor,图的一部分就可能无法回收。常见错误是把每一步的 loss 直接 append 到列表,导致整个历史图被保留:
loss_history = []
for batch in batches:
loss = model(batch).mean()
loss.backward()
loss_history.append(loss.item()) # 只保存数字,推荐
# loss_history.append(loss) # 可能保留图,不推荐
日志、指标和进度条只需要标量数字时,用 .item() 或 float(loss.detach())。但在 CUDA 上,.item() 可能触发设备同步;高吞吐训练不要在每个小算子后频繁取标量,可以按间隔批量记录。若要保留预测用于后处理,用 output.detach(),并根据需要 .cpu();不要把带图 GPU Tensor 长期放在 Python 容器里。
第 38 篇讲过 device 边界:detach() 不会自动把数据搬回 CPU,.cpu() 也不会自动断图。两件事都需要时写成 output.detach().cpu(),顺序表达“先停止梯度,再跨设备取结果”。
AI Infra 视角:训练显存大致由参数、梯度、优化器状态、激活和临时 workspace 共同构成。autograd 主要影响激活保存和反向工作区;减少 batch、开启检查点、使用混合精度或关闭不必要的输入梯度,都是在改变这条内存账本。推理服务则应彻底切断训练图,避免一次误用让每个请求都积累激活。

图4 计算图与显存生命周期
五、一个可检查的最小训练闭环
5.1 从 device 选择到参数梯度
把第 38 篇的 device 选择和本篇的 autograd 放在一起,可以得到一个最小闭环。这里先用一个手工参数,避免把 nn.Module 的细节提前到第 40 篇:
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.tensor([[1.0], [2.0], [3.0]], device=device)
target = 2 * x + 1
w = torch.tensor([[0.0]], device=device, requires_grad=True)
b = torch.tensor([0.0], device=device, requires_grad=True)
for step in range(3):
prediction = x @ w + b
loss = (prediction - target).pow(2).mean()
loss.backward()
with torch.no_grad():
w -= 0.1 * w.grad
b -= 0.1 * b.grad
w.grad = None
b.grad = None
print(step, loss.item())
注意上面代码中 device、x 等行不要有多余缩进;为便于排版,正确可运行版本如下:
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.tensor([[1.0], [2.0], [3.0]], device=device)
target = 2 * x + 1
w = torch.tensor([[0.0]], device=device, requires_grad=True)
b = torch.tensor([0.0], device=device, requires_grad=True)
for step in range(3):
prediction = x @ w + b
loss = (prediction - target).pow(2).mean()
loss.backward()
with torch.no_grad():
w -= 0.1 * w.grad
b -= 0.1 * b.grad
w.grad = None
b.grad = None
print(step, loss.item())
手工更新时,no_grad() 很重要:参数更新不是要学习的前向操作,若让它进入图,下一轮会把更新历史串起来,导致图不断变大或出现原地修改错误。第 42 篇会用优化器封装这个流程,但底层仍然是“读梯度、更新参数、清空梯度”。
5.2 用 autograd.grad 做局部检查
有时不想把梯度写入 .grad,只想查询某个输出对某个输入的导数,可以使用 torch.autograd.grad:
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x
grad = torch.autograd.grad(y, x)
print(grad) # (tensor(14.),),因为 3*x**2 + 2 = 14
它适合梯度检查、元学习和只需要局部梯度的代码。默认情况下,它也会释放图;若后面还需要二阶梯度,使用 create_graph=True。autograd.grad 返回梯度,不会像 backward() 那样自动累积到参数 .grad,两者不要混用后又期待同一个缓冲区自动更新。
5.3 梯度检查的工程方法
当自定义算子或复杂 loss 出现异常,先缩小问题:使用 CPU、极小 Tensor、固定随机种子和 float64;打印 requires_grad、is_leaf、grad_fn、shape、device 与 dtype。再用有限差分近似:对参数加一个很小的 epsilon,比较损失变化与梯度预测。有限差分不适合大模型训练,但适合验证局部公式。
也要检查数值本身。梯度为 nan 或 inf 可能来自除零、log(0)、指数溢出、损失尺度过大或混合精度下的下溢。不要看到 backward() 成功就认为训练正确:图能走通只说明代码在执行,不能证明梯度方向和数据归一化符合预期。

图5 AI Infra 中的训练反向路径
六、把 autograd 放回 AI Infra
6.1 训练和推理是两种完全不同的执行合同
训练请求需要保存激活、计算梯度、更新参数,通常是计算密集且显存密集的三段流水线;推理请求只需要前向和输出后处理,不应构建反向图。一个服务如果忘记使用 inference_mode(),小模型可能暂时看不出问题,大模型却会出现更高显存、更多 allocator 压力和更长尾延迟。
在 GPU 服务中,autograd 还和异步执行相互作用。CUDA kernel 可能异步排队,而 .item()、把结果转成 NumPy、某些跨 device 拷贝会触发同步。测量训练或推理延迟时,要区分前向、反向、数据搬运和同步边界;不要只用 Python 墙上时间就断言某个算子更快。第 38 篇的 device 检查在这里继续生效:图上的每个 Tensor 仍必须位于兼容 device。
6.2 分布式训练中的梯度不是只存在本地
单卡时,参数的 .grad 看起来像本地属性;数据并行训练中,每张卡先对自己的 micro-batch 做反向,然后通过通信原语聚合梯度。梯度同步的时机、bucket 大小、通信与反向的重叠,会直接影响吞吐。于是“什么时候调用 backward”不仅是数学问题,也是通信调度问题。
梯度累积可以减少同步频率:多个 micro-batch 的反向在本地累积,达到一个逻辑 batch 后再同步或更新。但必须明确损失缩放、同步上下文和显存上限。AI Infra 工程师在看 profiler 时,既要看反向 kernel,也要看 all-reduce 是否与反向重叠、是否因为某个小参数 bucket 造成气泡。
6.3 激活检查点与自动求导的交换
激活检查点(activation checkpointing)不保存某些中间激活,反向时重新执行对应前向,以时间换显存。它改变的是 autograd 的保存策略,不改变最终梯度的数学目标。模型越深、序列越长,激活越可能成为训练显存瓶颈;但重算会增加计算量,且随机算子、缓存和副作用必须谨慎处理。
混合精度、梯度缩放、编译器和算子融合也都围绕同一个图工作:前向图决定要保存什么,反向图决定如何生成梯度,运行时再把它们映射到 CUDA kernel 或融合算子。理解基础 autograd,才能看懂为什么 profiler 中会出现一串反向 kernel,为什么某个自定义 Python 操作阻碍融合,或者为什么一个不必要的 clone() 让显存峰值上升。
七、踩坑大全
坑 1:loss.backward() 报“没有 grad_fn”。 原因通常是前向被 no_grad() 或 inference_mode() 包住,或者在中途使用了 detach()、.item() 把 Tensor 变成了数值。检查 loss.requires_grad 和 loss.grad_fn,确认训练路径没有被断开。
坑 2:x.grad 是 None。 x 可能不是叶子 Tensor,或者它没有 requires_grad=True,也可能在反向前从未参与 loss。对中间 Tensor 调用 retain_grad(),对输入明确调用 requires_grad_(),并逐段检查图。
坑 3:梯度越来越大或每一步结果不稳定。 .grad 默认累积,训练循环却没有清零。把 zero_grad(set_to_none=True) 或显式的 param.grad = None 放在正确的更新边界;做梯度累积时则按设计保留并正确缩放 loss。
坑 4:第二次反向时报图已释放。 默认 backward 用完缓存就释放。优先重新前向或合并 loss;只有确实需要复用同一图时才用 retain_graph=True,并监控显存。
坑 5:把 retain_graph=True 和 create_graph=True 混为一谈。 前者是保留已有图以便再次反向,后者是把一阶梯度纳入新图以求高阶导。错误使用二者会把内存开销放大。
坑 6:推理显存持续上涨。 可能把带图的输出或 loss 存在列表里。日志保存 .item(),缓存预测保存 detach(),纯推理入口使用 model.eval() 加 inference_mode()。
坑 7:以为 detach() 会复制或搬到 CPU。 detach() 通常共享存储且仍在原 device;需要独立副本用 detach().clone(),需要交给 NumPy 用 detach().cpu().numpy()。
坑 8:eval() 不能关闭 autograd。 eval() 只切换模块行为,不负责图记录。验证和推理要按需求额外使用 no_grad() 或 inference_mode()。
坑 9:非标量输出直接 backward()。 向量没有唯一的反向起点。先 .sum()、.mean() 得到标量,或传入同形状的梯度种子。
坑 10:手工更新参数时没有 no_grad()。 更新表达式被记录进下一张图,容易产生图链、版本计数错误或意外梯度。参数更新属于优化器控制面,不应成为模型前向图的一部分。
坑 11:CPU 示例搬到 GPU 后速度反而更慢。 可能每步都在 CPU/GPU 之间来回拷贝,或在循环中频繁 .item() 触发同步。把连续算子留在同一 device,批量取结果,再用 CUDA event 或同步后的时间做基准。
坑 12:把图能反向当作梯度正确。 ReLU 死区、数值溢出、损失缩放和标签错误都可能让结果无效。用小规模手算、有限差分、异常梯度检查和可视化共同验证。
八、小结
- autograd 把前向操作组织成动态计算图,反向按链式法则计算损失对叶子 Tensor 的梯度。
requires_grad决定是否追踪路径;叶子 Tensor 的.grad是训练状态,非叶子要用retain_grad()才会保留。backward()默认要求标量,并且会把梯度累加到已有缓冲区;清零和更新是训练循环的必需步骤。detach()断开一个 Tensor 的图,no_grad()关闭代码块追踪,inference_mode()为纯推理提供更严格、更轻量的模式;它们都不等同于eval()。retain_graph用于复用已有图,create_graph用于构建高阶梯度图;两者都会影响显存,不能当作通用修复。- 从 AI Infra 角度,autograd 决定训练激活、反向 kernel、梯度通信和显存生命周期;推理则应明确关闭训练图并测量同步边界。
九、课后练习
- (基础) 创建一个
requires_grad=True的标量x,计算y = 3*x**2 + 2*x + 1,调用backward()并手算验证梯度。 - (基础) 构造长度为 3 的向量输出,分别用
.sum()、.mean()和自定义梯度种子调用反向,比较三种.grad。 - (理解) 让一个中间 Tensor 调用
retain_grad(),打印它和输入的梯度,解释为什么两者不同。 - (理解) 写一个两步训练循环,先故意不清零梯度,再加入
grad = None,观察累积差异。 - (理解) 分别用
detach()、detach().clone()和no_grad()改写一个教师—学生输出示例,说明共享存储和断图边界。 - (进阶) 用
create_graph=True求x**4的一阶和二阶导,并估算它比一阶反向多保存了什么。 - (进阶) 对一个小型线性回归做有限差分梯度检查,分别在 CPU 和第 38 篇选择的
device上运行。 - (AI Infra) 设计一个推理函数,记录普通前向、
no_grad()和inference_mode()的显存与耗时;说明为什么基准中要处理 CUDA 同步。
下一篇预告
第40篇|nn.Module 与参数管理。本篇用手工 Tensor 完成了“前向—反向—更新”的骨架,下一篇会把它组织成可复用的 nn.Module:理解参数注册、子模块树、state_dict、train()/eval() 与 .to(device) 如何协同,让第 38 篇的设备迁移和本篇的 autograd 真正进入模型工程。再往后,我们会用 Dataset、DataLoader 和训练循环把单步计算扩展为可观测、可恢复的训练作业,最终连接到优化器、混合精度和大模型服务。
