Shortcuts

轻松入门 torch.autograd

创建于:2017年3月24日 | 最后更新:2024年12月12日 | 最后验证:2024年11月5日

torch.autograd 是 PyTorch 的自动微分引擎,它为神经网络训练提供动力。在本节中,您将概念性地理解 autograd 如何帮助神经网络训练。

背景

神经网络(NNs)是一组嵌套函数,这些函数在某些输入数据上执行。这些函数由参数(包括权重和偏置)定义,在PyTorch中,这些参数存储在张量中。

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

前向传播: 在前向传播中,神经网络会对其正确的输出做出最佳猜测。它通过每个函数运行输入数据以做出这个猜测。

反向传播:在反向传播中,神经网络根据其猜测中的误差按比例调整其参数。它通过从输出向后遍历,收集误差相对于函数参数的导数(梯度),并使用梯度下降优化参数来实现这一点。有关反向传播的更详细讲解,请查看这个来自3Blue1Brown的视频

在PyTorch中的使用

让我们来看一个训练步骤的例子。 在这个例子中,我们从torchvision加载一个预训练的resnet18模型。 我们创建一个随机的数据张量来表示一个具有3个通道、高度和宽度为64的单张图像, 以及其对应的label,初始化为一些随机值。预训练模型中的标签形状为(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)
Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth

  0%|          | 0.00/44.7M [00:00<?, ?B/s]
 46%|####6     | 20.8M/44.7M [00:00<00:00, 217MB/s]
 93%|#########3| 41.6M/44.7M [00:00<00:00, 218MB/s]
100%|##########| 44.7M/44.7M [00:00<00:00, 218MB/s]

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

prediction = model(data) # forward pass

我们使用模型的预测和相应的标签来计算误差(loss)。 下一步是通过网络反向传播这个误差。 当我们调用误差张量的.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是如何收集梯度的。我们创建了两个张量ab,并将requires_grad=True。这向autograd发出信号,表示应该跟踪它们的每一个操作。

import torch

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

我们从ab创建另一个张量Q

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

假设ab是神经网络的参数,Q是误差。在神经网络训练中,我们需要误差相对于参数的梯度,即。

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

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

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

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

同样地,我们也可以将 Q 聚合成一个标量并隐式调用 backward,如 Q.sum().backward()

梯度现在存储在 a.gradb.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} = m \cdot \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) = m \cdot \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 在一个由 Function 对象组成的有向无环图(DAG)中记录数据(张量)和所有执行的操作(以及生成的新张量)。在这个 DAG 中,叶子节点是输入张量,根节点是输出张量。通过从根节点到叶子节点追踪这个图,你可以使用链式法则自动计算梯度。

在前向传播过程中,autograd 同时做两件事:

  • 运行请求的操作以计算结果张量,并且

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

当在DAG根上调用.backward()时,反向传播开始。autograd然后:

  • 计算每个.grad_fn的梯度,

  • 将它们累积在相应张量的.grad属性中,并且

  • 使用链式法则,一直传播到叶张量。

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

../../_images/dag_autograd.png

注意

在PyTorch中,DAGs是动态的 需要注意的是,图是从头开始重新创建的;每次 .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

在神经网络中,不计算梯度的参数通常被称为冻结参数。 如果你提前知道不需要这些参数的梯度,那么“冻结”模型的一部分是有用的 (这通过减少自动梯度计算提供了一些性能优势)。

在微调过程中,我们冻结了模型的大部分部分,通常只修改分类器层以对新标签进行预测。 让我们通过一个小例子来演示这一点。和之前一样,我们加载一个预训练的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()中作为上下文管理器可用。


优云智算