torch.autograd 轻松入门

torch.autograd 轻松入门

torch.autograd 是 PyTorch 的自动微分引擎，为神经网络训练提供核心支持。本节将帮助你从概念层面理解 autograd 如何助力神经网络训练。

背景知识

神经网络（NN）是由嵌套函数组成的集合，这些函数作用于输入数据并产生输出。这些函数由参数（包含权重和偏置）定义，在 PyTorch 中这些参数存储在张量中。

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

1. 前向传播（Forward Propagation）：在前向传播阶段，神经网络根据输入数据尽可能做出最接近真实值的输出猜测。它会将输入数据依次通过网络中的每个函数来生成这个猜测结果。

2. 反向传播（Backward Propagation）：在反向传播阶段，神经网络会根据猜测结果的误差成比例地调整自身参数。具体过程是从输出端反向遍历网络，收集误差相对于各函数参数的导数（梯度），并使用梯度下降算法优化这些参数。想要更深入了解反向传播的原理，可以观看 3Blue1Brown 的相关视频。

PyTorch 中的使用示例

我们来演示单个训练步骤的完整流程。本示例中，我们从 torchvision 加载预训练的 resnet18 模型，创建一个随机数据张量（模拟单张 3 通道、64x64 尺寸的图像），并初始化对应的标签为随机值（预训练模型的标签形状为 (1,1000)）。

注意
本教程仅适用于 CPU 环境，无法在 GPU 设备上运行（即使将张量移至 CUDA 设备也不行）。

import torch
from torchvision.models import resnet18, ResNet18_Weights

# 加载预训练的 resnet18 模型
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]
 63%|██████▎   | 28.0M/44.7M [00:00<00:00, 293MB/s]
100%|██████████| 44.7M/44.7M [00:00<00:00, 280MB/s]

接下来，将输入数据通过模型的各层网络以生成预测结果，这就是前向传播过程：

prediction = model(data) # 前向传播

利用模型的预测结果和对应的标签计算误差（损失值）。下一步是将这个误差反向传播通过整个网络——只需在误差张量上调用 .backward() 方法即可触发反向传播。autograd 会自动计算每个模型参数的梯度，并将其存储在参数的 .grad 属性中。

loss = (prediction - labels).sum()
loss.backward() # 反向传播

然后加载优化器（此处使用学习率为 0.01、动量为 0.9 的 SGD），并将模型的所有参数注册到优化器中：

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

最后调用 .step() 方法执行梯度下降，优化器会根据存储在 .grad 中的梯度值调整每个参数：

optim.step() # 梯度下降

至此，你已经掌握了训练神经网络所需的核心流程。以下部分详细讲解 autograd 的工作原理，你可以根据需要选择阅读。

Autograd 中的微分计算

我们来具体看看 autograd 是如何收集梯度的。创建两个设置 requires_grad=True 的张量 a 和 b，这会告诉 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 参数。该参数是与 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：

# 验证收集到的梯度是否正确
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$：