PyTorch:自动求导模块

发布时间:2026-06-03 00:03  浏览量:1

PyTorch 的自动求导模块 torch.autograd,是深度学习训练中最核心的机制之一。它负责根据模型的前向计算过程,自动构建计算图,并在反向传播时自动计算模型参数的梯度。

简单地说,自动求导模块回答的是:模型输出如何受到参数影响,损失函数如何对参数求导,以及优化器为什么能够根据梯度更新模型参数。

如果说 torch.nn 模块负责组织神经网络结构,torch.optim 模块负责更新模型参数,那么 torch.autograd 模块就负责连接二者:它根据损失函数自动计算梯度,让模型能够通过训练逐步改进。

一、认识自动求导模块

自动求导(Automatic Differentiation)是 PyTorch 训练神经网络的基础能力。它并不是简单地对某个数学公式手工求导,而是在张量计算过程中自动记录操作历史,并在需要时沿着计算图反向传播梯度。

图 1:自动求导模块在 PyTorch 训练流程中的位置

在 PyTorch 中,一个典型训练流程通常包括:

• 构造输入数据

• 定义模型

• 执行前向传播

• 计算损失函数

• 执行反向传播

• 使用优化器更新参数

其中,自动求导主要发生在“前向传播之后、参数更新之前”。模型先根据输入得到预测结果,损失函数再衡量预测结果与真实目标之间的差异,随后 PyTorch 根据这条计算路径自动计算梯度。

一个最简单的例子如下:

输出结果为:

tensor(7.)

这里的数学关系是:

当 x = 2 时:

因此,x.grad 的结果是 7。

这个例子说明:只要张量启用了梯度跟踪,PyTorch 就可以根据前向计算过程自动完成反向求导。

二、requires_grad:是否跟踪梯度

在 PyTorch 中,并不是所有张量都会自动参与求导。只有设置了 requires_grad=True 的张量,才会被自动求导系统跟踪。

例如:

这里可以看出:

• x 需要计算梯度

• y 不需要计算梯度

• z 由 x 参与计算得到,因此 z 也会被纳入计算图

requires_grad 的作用可以理解为:告诉 PyTorch 是否需要记录这个张量相关的计算历史。

在神经网络中,模型参数通常默认需要计算梯度。例如:

bias True

这说明 nn.Linear 中的权重和偏置默认会参与梯度计算。训练时,优化器正是根据这些参数的梯度来更新它们。

三、计算图:记录张量之间的计算关系

自动求导的基础是(Computation Graph)。在前向计算过程中,PyTorch 会记录张量之间的依赖关系,并形成一张动态计算图。

例如:

这里的计算过程可以理解为:

→ y = b²

只要最终对 y 调用 backward,PyTorch 就会沿着这条路径反向计算梯度。

图 2:从前向计算到反向传播的计算图

需要注意的是,PyTorch 使用的是动态计算图。也就是说,计算图是在每次前向传播时动态构建的,而不是提前固定好的。

这使得 PyTorch 代码更接近普通 Python 程序,可以自然地使用条件判断、循环和函数调用。

例如:

这里的计算路径取决于运行时的条件判断。PyTorch 会根据实际执行过的操作构建计算图。

四、backward:执行反向传播

backward 是自动求导中最常用的方法。它的作用是从某个结果张量出发,沿着计算图反向计算梯度。

例如:

print(x.grad)

输出结果为:

tensor(6.)

对应的数学关系是:

当 x = 3 时,梯度为 6。

在神经网络训练中,通常不是直接对模型输出调用 backward,而是对损失函数调用:

loss.backward

这表示:计算损失函数对所有可学习参数的梯度。

例如:

这里可以理解为:

• 前向传播得到预测值 pred

• 损失函数得到 loss

• loss.backward 计算 loss 对模型参数的梯度

• 梯度保存在参数的 .grad 属性中

五、grad:查看张量的梯度

当一个张量需要计算梯度,并且参与了反向传播之后,它的梯度通常保存在 .grad 属性中。

例如:

print(x.grad)

输出结果为:

tensor(12.)

对应的数学关系是:

当 x = 2 时:

在神经网络中,也可以查看模型参数的梯度:

需要注意的是,通常只有叶子张量的 .grad 会被保留下来。模型参数一般就是叶子张量,因此它们的 .grad 可以直接查看。

中间计算结果虽然参与计算图,但默认不一定保留 .grad。如果确实需要查看中间张量的梯度,可以使用 retain_grad。

六、梯度累积与 zero_grad

PyTorch 中的梯度默认会累积,而不是每次自动清零。

例如:

第二次打印的结果不是只包含 y2 对 x 的梯度,而是两次梯度的累加。

这是因为 PyTorch 允许在复杂场景中多次累积梯度,例如梯度累积训练、多个损失共同反向传播等。但在普通训练循环中,每一轮参数更新前通常都需要清空上一轮梯度。

因此,训练代码中经常会写:

optimizer.step

完整示例如下:

这个顺序非常重要。一般训练时不应忘记 optimizer.zero_grad,否则梯度会不断累积,导致参数更新异常。

七、标量输出与非标量输出的 backward

在 PyTorch 中,如果输出是标量,通常可以直接调用 backward。

例如:

这里的 y 是一个标量,因此可以直接反向传播。

但如果输出不是标量,就不能直接调用 backward,除非提供一个与输出形状一致的梯度参数。

例如:

这里的 torch.ones_like(y) 可以理解为给每个输出分量提供一个反向传播的初始权重。

在多数神经网络训练中,损失函数通常会把多个样本的损失聚合为一个标量,所以常见写法仍然是:

loss.backward

这也是为什么损失函数默认通常会返回一个标量结果。

八、torch.autograd.grad:直接计算梯度

除了使用 backward,PyTorch 还提供了 torch.autograd.grad,用于直接计算某些输出对某些输入的梯度。

例如:

(tensor(14.),)

对应的数学关系是:

当 x = 2 时:

backward 和 autograd.grad 的区别可以简单理解为:

• backward 通常把梯度累积到张量的 .grad 属性中

• autograd.grad 通常直接返回梯度结果

在普通神经网络训练中,loss.backward 更常见;在需要精确控制梯度计算的场景中,torch.autograd.grad 更灵活。例如,一些自定义优化方法、元学习、高阶导数计算和物理信息神经网络中,会经常用到它。

九、关闭梯度计算:no_grad 与 inference_mode

并不是所有前向计算都需要梯度。在验证、测试和推理阶段,模型通常只需要输出结果,不需要反向传播。

这时可以使用 torch.no_grad 关闭梯度记录:

这样做有两个主要好处:

• 减少不必要的计算图记录

• 降低内存占用

在推理阶段,常见写法是:

还可以使用 torch.inference_mode:

inference_mode 比 no_grad 更偏向纯推理场景,限制更强,通常可以进一步减少自动求导相关开销。初学阶段可以先掌握 no_grad,在明确只做推理时再了解 inference_mode。

需要注意的是:

• model.eval 负责切换模型中 Dropout、BatchNorm 等层的行为

• torch.no_grad 负责关闭梯度记录

二者作用不同,推理时通常一起使用。

十、detach:从计算图中分离张量

detach 用于从当前计算图中分离出一个张量。分离后的张量与原张量共享数据,但不再保留梯度计算关系。

例如:

输出结果为:

False

这说明 z 不再被自动求导系统跟踪。

detach 常用于以下场景:

• 在记录日志时保存数值,但不希望继续跟踪梯度

• 在某些模型结构中阻断梯度传播

• 在生成模型或强化学习中分离目标值

• 把张量转换为 NumPy 数组前先断开计算图

例如:

value = loss.detach.item

或者:

array = tensor.detach.cpu.numpy

需要注意的是,detach 不是复制一份完全独立的数据。它返回的新张量可能与原张量共享底层存储,因此不应随意对其进行原地修改。

十一、冻结参数:控制哪些部分参与训练

在迁移学习或微调模型时,经常需要冻结部分参数,只训练新加的分类头或少量模块。

冻结参数的关键是设置:

param.requires_grad = False

例如:

这里可以理解为:

• feature 部分参数不再计算梯度

• classifier 部分参数仍然参与训练

• 优化器只接收需要更新的参数

这种做法常用于预训练模型微调。例如,先冻结主干网络,只训练任务相关的输出层;之后再根据需要解冻部分层进行进一步训练。

十二、原地操作与计算图问题

在 PyTorch 中,某些原地操作可能破坏自动求导所需的中间值,从而导致反向传播错误。

原地操作通常指会直接修改原张量内容的操作,例如:

x.relu_

带下划线的方法通常表示原地修改。

并不是所有原地操作都会出错,但在涉及梯度计算时,需要格外谨慎。因为自动求导系统在反向传播时可能需要前向传播中的某些中间结果,如果这些结果被原地修改,梯度计算就可能无法正确进行。

例如,下面这种写法就容易引发问题:

在实际训练中,建议初学者少用原地操作,尤其是在不确定其影响时。优先使用非原地写法:

y = y + 1

而不是:

y.add_(1)

这样更安全,也更容易调试。

十三、自动求导与神经网络训练闭环

自动求导并不是孤立存在的,它通常与模型、损失函数和优化器共同构成训练闭环。

图 3:自动求导驱动的神经网络训练闭环

下面是一个完整的简单分类训练示例:

这个示例中,自动求导主要体现在:

• 前向传播时,PyTorch 自动记录计算路径

• loss.backward 自动计算损失对参数的梯度

• 梯度保存在模型参数的 .grad 属性中

• optimizer.step 根据梯度更新参数

这就是深度学习训练最基本的工作方式。

十四、验证与推理阶段的自动求导控制

训练阶段需要梯度,验证和推理阶段通常不需要梯度。因此,训练和推理代码应当区分清楚。

训练阶段通常写成:

验证或推理阶段通常写成:

两者的核心区别是:

• 训练阶段需要记录计算图

• 训练阶段需要反向传播

• 训练阶段需要更新参数

• 推理阶段只需要前向计算

• 推理阶段通常不需要梯度

如果在推理阶段忘记使用 torch.no_grad,代码通常仍然能运行,但会记录不必要的计算图,增加内存开销。

如果在训练阶段错误使用了 torch.no_grad,则可能导致损失无法反向传播,模型参数无法更新。

十五、使用自动求导时应注意的问题

1、只有 requires_grad=True 的张量才会被跟踪

普通张量默认不计算梯度。例如:

x = torch.tensor(2.0)

如果需要对 x 求导,应写成:

x = torch.tensor(2.0, requires_grad=True)

在神经网络中,模型参数默认通常已经设置了 requires_grad=True,不需要手动指定。

2、通常对标量 loss 调用 backward

普通训练中,推荐对损失函数结果调用:

loss.backward

如果输出不是标量,则需要提供额外的梯度参数。初学阶段可以先理解:损失函数通常会把多个样本的误差汇总成一个标量,方便执行反向传播。

3、反向传播前要清空旧梯度

PyTorch 的梯度默认会累积,所以训练循环中通常要写:

如果忘记清空梯度,参数更新可能会受到前几轮梯度的影响,导致训练过程异常。

4、不要把 loss.item 用于反向传播

loss.item 会把张量转换为普通 Python 数值,它已经脱离计算图。

正确写法是:

loss.backward

错误写法是:

loss.item.backward

item 适合用于打印日志,而不是用于继续计算梯度。

5、不要在训练阶段随意使用 detach

detach 会切断计算图。如果在训练过程中错误地对关键张量使用 detach,梯度就无法继续传播。

例如:

loss.backward

这种写法会导致模型参数无法通过 logits 接收到梯度,训练失效。

6、推理阶段应使用 no_grad

验证或推理时推荐写法:

output = model(x)

这样可以避免记录不必要的计算图,减少内存消耗。

7、冻结参数要同时考虑优化器

如果某些参数设置了:

这样可以让训练目标更明确,也避免优化器持有不需要更新的参数。

8、注意原地操作可能影响梯度计算

带下划线的方法通常是原地操作,例如:

zero_

在不熟悉自动求导机制时,应尽量少在参与计算图的张量上使用原地操作。否则可能破坏反向传播所需的中间结果。

小结

PyTorch 自动求导模块负责记录计算图并自动计算梯度,是模型训练的核心机制。理解 requires_grad、backward、.grad、梯度清零与推理阶段关闭梯度,有助于把握神经网络如何从损失中学习并更新参数。

“点赞有美意,赞赏是鼓励”