【PyTorch单点知识】自动求导机制的原理与实践

0. 前言

按照国际惯例，首先声明：本文只是我自己学习的理解，虽然参考了他人的宝贵见解及成果，但是内容可能存在不准确的地方。如果发现文中错误，希望批评指正，共同进步。

在深度学习中，自动求导（Automatic Differentiation, AD）是一项至关重要的技术，它使我们能够高效地计算神经网络的梯度，进而通过反向传播算法更新权重。

PyTorch作为一款动态计算图的深度学习框架，以其灵活性和易用性著称，其自动求导机制是其实现高效、灵活训练的核心。本文将深入探讨PyTorch中的自动求导机制，从原理到实践，通过代码示例来展示其工作流程。

如果对计算图不太了解，可以参考我的往期文章：基于TorchViz详解计算图（附代码）

1. 自动求导的基本原理

自动求导是一种数学方法，用于计算函数的导数。与数值微分相比，自动求导能够提供精确的导数计算结果，同时避免了符号微分中可能出现的手动求导错误。在深度学习中，我们通常关注的是反向模式backward的自动求导，即从输出向输入方向传播梯度的过程。

反向模式自动求导基于链式法则，它允许我们将复杂的复合函数的导数分解成多个简单函数的导数的乘积。在神经网络中，每一层都可以看作是一个简单的函数，通过链式法则，我们可以从前向传播的输出开始，逆向计算每个参数的梯度。

2. PyTorch中的自动求导

PyTorch通过其autograd模块实现了自动求导机制。autograd记录了所有的计算步骤，创建了一个计算图（Computational Graph），并在需要时执行反向传播，计算梯度。

2.1 创建计算图

在PyTorch中，当一个张量（Tensor）的requires_grad=True时，任何对该张量的操作都会被记录在计算图中。例如：

import torch

x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

print(y.grad_fn)  # 查看y的计算节点
print(z.grad_fn)  # 查看z的计算节点

输出为：

<AddBackward0 object at 0x000001CADEC6AB60>
<MulBackward0 object at 0x000001CADEC6AB60>

在上述代码中，z的计算节点显示了z是如何由y计算得来的，而y的计算节点则显示了y是如何由x计算得来的。这样就形成了一个计算图。

2.2 反向传播

一旦我们完成了前向传播并得到了最终的输出，就可以调用out.backward()来进行反向传播，计算梯度。例如：

import torch

x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

out.backward()
print(x.grad)

这里，x.grad就是out相对于x的梯度。输出为：

tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])

2.3 反向传播详解

下面我们来详细分析下1.2节的具体计算过程：

1. 首先，创建了一个2x2的张量x，其值全为1，并且设置了requires_grad=True，这意味着PyTorch将会追踪这个张量上的所有操作，以便能够计算梯度。

x = torch.ones(2, 2, requires_grad=True)

2. 然后，将x与2相加得到y：

y = x + 2

此时y的值为：

tensor([[3., 3.],
        [3., 3.]])

3. 接下来，将y的每个元素平方再乘以3得到z：

z = y * y * 3

此时z的值为：

tensor([[27., 27.],
        [27., 27.]])

4. 计算z的平均值作为输出out：

out = z.mean()

此时out的值为：

tensor(27.)

5. 使用backward()函数对out进行反向传播，计算梯度：

out.backward()

6. 最后，打印x的梯度：

print(x.grad)

由于out是通过一系列操作从x得到的，我们可以根据链式法则计算出x的梯度。具体来说，out相对于x的梯度可以通过以下步骤计算得出：

• out相对于z的梯度是1/z.size(0)（因为z.mean()是对z的所有元素取平均），这里z.size(0)等于4，所以out相对于z的梯度是1/4。
• z相对于y的梯度是y * 3 * 2（因为z = y^2 * 3，所以dz/dy = 2*y*3）。
• y相对于x的梯度是1（因为y = x + 2，所以dy/dx = 1）。

综合以上，out相对于x的梯度是：

1/4 * (y * 3 * 2) * 1

由于y的值为[[3, 3], [3, 3]]，那么上述梯度计算结果为：

1/4 * (3 * 3 * 2) * 1 = 9/2 = 4.5

因此，最终x.grad的值为：

tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])

2.4 梯度清零

在多次迭代中，梯度会累积在张量中，因此在每次迭代开始之前，我们需要调用optimizer.zero_grad()来清零梯度，防止梯度累积。（PyTorch为了训练方便，会默认梯度累积）

2.5 定制自动求导

PyTorch还允许我们定义自己的自动求导函数，通过继承torch.autograd.Function类并重写forward和backward方法。这为实现更复杂的计算提供了可能。

3. 代码实例：线性回归的自动求导

接下来，我们将通过一个简单的线性回归问题，演示PyTorch自动求导机制的实际应用。

假设我们有一组数据点，我们想找到一条直线(y = wx + b)，使得这条直线尽可能接近这些数据点。我们的目标是最小化损失函数（例如均方误差）。

import torch
import numpy as np

# 准备数据
np.random.seed(0)
X = np.random.rand(100, 1)
Y = 2 + 3 * X + 0.1 * np.random.randn(100, 1)

X = torch.from_numpy(X).float()
Y = torch.from_numpy(Y).float()

# 初始化权重和偏置
w = torch.tensor([1.], requires_grad=True)
b = torch.tensor([1.], requires_grad=True)

# 定义模型和损失函数
def forward(x):
    return w * x + b

loss_fn = torch.nn.MSELoss()

# 训练循环
learning_rate = 0.01
for epoch in range(1000):
    # 前向传播
    y_pred = forward(X)
    
    # 计算损失
    loss = loss_fn(y_pred, Y)
    
    # 反向传播
    loss.backward()
    
    # 更新权重
    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad
        
        # 清零梯度
        w.grad.zero_()
        b.grad.zero_()
    
    if (epoch+1) % 100 == 0:
        print(f'Epoch [{epoch+1}/1000], Loss: {loss.item():.4f}')

print('Final weights:', w.item(), 'bias:', b.item())

输出：

Epoch [100/1000], Loss: 0.1273
Epoch [200/1000], Loss: 0.0782
Epoch [300/1000], Loss: 0.0620
Epoch [400/1000], Loss: 0.0497
Epoch [500/1000], Loss: 0.0404
Epoch [600/1000], Loss: 0.0332
Epoch [700/1000], Loss: 0.0277
Epoch [800/1000], Loss: 0.0235
Epoch [900/1000], Loss: 0.0203
Epoch [1000/1000], Loss: 0.0179
Final weights: 2.68684983253479 bias: 2.17771577835083

在这个例子中，我们首先准备了一些随机生成的数据，然后初始化了权重w和偏置b。在训练循环中，我们通过前向传播计算预测值，使用均方误差损失函数计算损失，然后通过调用loss.backward()进行反向传播，最后更新权重和偏置。通过多次迭代，我们最终找到了使损失最小化的权重和偏置。

4. 结论

PyTorch的自动求导机制是其强大功能的关键所在。通过autograd模块，PyTorch能够自动跟踪计算图并高效地计算梯度，这大大简化了深度学习模型的开发过程。本文通过理论解释和代码示例，深入探讨了PyTorch中的自动求导机制，希望读者能够从中获得对这一重要概念的深刻理解，并在实际项目中灵活运用。