torch.autograd 的基础介绍

Created On: Mar 24, 2017 | Last Updated: Jan 10, 2025 | Last Verified: Nov 05, 2024

torch.autograd 是 PyTorch 的自动微分引擎，为神经网络训练提供支持。在本节中，您将从概念上了解 autograd 如何帮助神经网络进行训练。

背景知识

神经网络 (NN) 是一系列嵌套函数的集合，这些函数在某些输入数据上执行操作。这些函数通过*参数*（包括权重和偏置）来定义，在 PyTorch 中，这些参数存储在张量中。

训练神经网络分为两个步骤：

前向传播：在前向传播中，神经网络对正确的输出进行最佳猜测。它通过每个函数运行输入数据以进行猜测。

反向传播：在反向传播中，神经网络根据错误大小调整其参数。它从输出开始反向遍历，收集关于函数参数的误差的导数（梯度），并使用梯度下降优化这些参数。有关反向传播更详细的讲解，请查看此 视频 (3Blue1Brown)。

在 PyTorch 中的使用

让我们看看单次训练步骤。对于此示例，我们从 torchvision 加载预训练的 resnet18 模型。我们创建一个随机张量数据，代表一张具有 3 个通道，64 高和 64 宽的图像，以及它的对应 标签，初始化为一些随机值。在预训练模型中，标签形状为 (1,1000)。

备注

本教程仅在 CPU 上使用，不在 GPU 上运行（即便将张量移动到 CUDA 也不支持）。

import torch
from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)
data = torch.rand(1, 3, 64, 64)
labels = torch.rand(1, 1000)

接下来，我们通过模型的每一层运行输入数据以进行预测。这是 前向传播。

prediction = model(data) # forward pass

我们使用模型的预测和对应的标签来计算误差（损失）。下一步是通过网络反向传播这个误差。当我们在误差张量上调用 .backward() 时，反向传播开始。Autograd 会计算并存储每个模型参数的梯度，在参数的 .grad 属性中。

loss = (prediction - labels).sum()
loss.backward() # backward pass

接下来，我们加载一个优化器，这里使用学习率为 0.01，动量 为 0.9 的 SGD。我们将模型的所有参数注册到该优化器中。

optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

最后，我们调用 .step() 来启动梯度下降。优化器会根据保存在 .grad 中的梯度调整每个参数。

optim.step() #gradient descent

到现在为止，您已经掌握了训练神经网络所需的一切。下面的部分详细说明了 autograd 的工作原理，您可以根据需要选择阅读或跳过。

---

在 Autograd 中的求导

让我们看看 autograd 如何收集梯度。我们创建了两个张量 a 和 b ，并设置 requires_grad=True。这告诉 autograd 应跟踪所有对它们的操作。

import torch

a = torch.tensor([2., 3.], requires_grad=True)
b = torch.tensor([6., 4.], requires_grad=True)

我们通过 a 和 b 创建了另一个张量 Q。

\[Q = 3a^3 - b^2 \]

Q = 3*a**3 - b**2

假设 a 和 b 是神经网络的参数，而 Q 是误差。在神经网络训练中，我们需要误差相对于参数的梯度，即：

\[\frac{\partial Q}{\partial a} = 9a^2 \]

\[\frac{\partial Q}{\partial b} = -2b \]

当我们在 Q 上调用 .backward() 时，autograd 会计算这些梯度并将其存储在相关张量的 .grad 属性中。

由于 Q 是一个向量，因此我们需要在 Q.backward() 中显式传递一个 gradient 参数。gradient 是一个与 Q 形状相同的张量，表示 Q 相对于自身的梯度，即：

\[\frac{dQ}{dQ} = 1 \]

或者，我们也可以将 Q 聚合为标量并隐式调用 backward，例如 Q.sum().backward()。

external_grad = torch.tensor([1., 1.])
Q.backward(gradient=external_grad)

现在，梯度会存放在 a.grad 和 b.grad 中。

# check if collected gradients are correct
print(9*a**2 == a.grad)
print(-2*b == b.grad)

tensor([True, True])
tensor([True, True])

可选阅读 - 使用 autograd 进行向量微积分

从数学上讲，如果您有一个向量值函数 \(\vec{y}=f(\vec{x})\)，那么 \(\vec{y}\) 对于 \(\vec{x}\) 的梯度是一个雅可比矩阵 \(J\)：

\[J = \left(\begin{array}{cc} \frac{\partial \bf{y}}{\partial x_{1}} & ... & \frac{\partial \bf{y}}{\partial x_{n}} \end{array}\right) = \left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{1}}{\partial x_{n}}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_{m}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)\]

总的来说，torch.autograd 是一个向量-雅可比积的计算引擎。即给定任何向量 \(\vec{v}\)，计算乘积 \(J^{T}\cdot \vec{v}\)。

如果 \(\vec{v}\) 恰好是标量函数 \(l=g\left(\vec{y}\right)\) 的梯度：

\[\vec{v} = \left(\begin{array}{ccc}\frac{\partial l}{\partial y_{1}} & \cdots & \frac{\partial l}{\partial y_{m}}\end{array}\right)^{T}\]

那么根据链式法则，向量-雅可比积将是 \(l\) 相对于 \(\vec{x}\) 的梯度：

\[J^{T}\cdot \vec{v} = \left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{1}}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_{1}}{\partial x_{n}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)\left(\begin{array}{c} \frac{\partial l}{\partial y_{1}}\\ \vdots\\ \frac{\partial l}{\partial y_{m}} \end{array}\right) = \left(\begin{array}{c} \frac{\partial l}{\partial x_{1}}\\ \vdots\\ \frac{\partial l}{\partial x_{n}} \end{array}\right)\]

向量-雅可比积的这种特性正是我们在上述示例中使用的；external_grad 表示 \(\vec{v}\)。

计算图

从概念上讲，autograd 会记录数据（张量）及所有执行的操作（以及产生的新张量），构成一个有向无环图 (DAG)，其中包含 Function 对象。在此图中，叶子是输入张量，根是输出张量。通过从根到叶子的这条路径，您可以使用链式法则自动计算梯度。

在前向传播时，autograd 同时执行两个操作：

• 运行请求操作以计算结果张量，

• 在 DAG 中维护操作的 梯度函数。

当在 DAG 根节点调用 .backward() 时，反向传播开始。然后 autograd：

• 从每个 .grad_fn 中计算梯度，

• 将梯度累加到相应张量的 .grad 属性中，

• 使用链式法则，一路传播到叶子张量。

以下是我们的示例中 DAG 的可视化表示。在图中，箭头表示前向传播的方向。节点代表前向传播中每个操作的反向函数。蓝色叶节点表示我们的叶张量 a 和 b。

备注

PyTorch 中的 DAG 是动态的。需要注意的一点是，在每次调用 .backward() 后，图都会从头开始重新创建；autograd 开始填充一个新的图。这正是使您可以在模型中使用控制流语句的原因；如果需要，您可以在每次迭代时更改形状、大小和操作。

从 DAG 中排除

torch.autograd 会跟踪所有 requires_grad 标志设置为 True 的张量上的操作。对于不需要梯度的张量，将此属性设置为 False 可将其从梯度计算 DAG 中排除。

如果仅有一个输入张量设置 requires_grad=True，操作的输出张量仍然需要计算梯度。

x = torch.rand(5, 5)
y = torch.rand(5, 5)
z = torch.rand((5, 5), requires_grad=True)

a = x + y
print(f"Does `a` require gradients?: {a.requires_grad}")
b = x + z
print(f"Does `b` require gradients?: {b.requires_grad}")

Does `a` require gradients?: False
Does `b` require gradients?: True

在神经网络中，不计算梯度的参数通常称为 冻结参数。如果您事先知道某些参数不需要梯度，可以“冻结”模型的一部分（通过减少 autograd 计算可提高性能）。

在微调中，我们冻结大部分模型，通常只修改分类层以针对新标签进行预测。让我们通过一个小示例来演示这一点。如前，我们加载了一个预训练的 resnet18 模型，并冻结了所有参数。

from torch import nn, optim

model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze all the parameters in the network
for param in model.parameters():
    param.requires_grad = False

假设我们想要在 10 个标签的新数据集上微调模型。在 resnet 中，分类器是最后的线性层 model.fc。我们只需用一个新的线性层（默认未冻结）替换它，该层充当我们的分类器。

model.fc = nn.Linear(512, 10)

现在，模型中除 model.fc 的参数外，所有参数都被冻结。唯一计算梯度的参数是 model.fc 的权重和偏置。

# Optimize only the classifier
optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

请注意，虽然我们将所有参数都注册到优化器中，但唯一计算梯度（并因此在梯度下降中更新）的参数是分类器的权重和偏置。

相同的排除功能可以作为上下文管理器 torch.no_grad() 中提供。

---

进一步阅读：

• 原地操作和多线程 Autograd

• 反向模式自动微分的示例实现

• 视频：深入教程：PyTorch Autograd 说明