torch中的梯度计算

1. 引言

梯度计算是深度学习的核心机制之一，PyTorch通过 自动微分（Autograd） 系统实现了高效的梯度计算。本文将深入解析PyTorch中梯度计算的原理与实现细节。

2. 计算图与反向传播

PyTorch采用 动态计算图（Dynamic Computation Graph）：

• 前向传播时实时构建计算图
• 每个张量操作被记录为图中的节点
• 反向传播时自动计算梯度

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3*x  # 构建计算图
y.backward()       # 自动反向传播
print(x.grad)      # 输出梯度值 7.0

3. 核心机制

3.1 Tensor属性

• requires_grad: 控制梯度跟踪开关,表示是否要对当前变量计算梯度，如果为True，则在该变量上的所有操作都会被追踪。
• grad: 存储计算得到的梯度值，这个实际存储的是，反向传播过程中，当前节点的梯度值。
• grad_fn: 记录创建该Tensor的操作，用于反向传播时计算梯度。这个变量只存储了，最后一个操作的信息，而不是所有的操作的信息。当我们需要得到前面的操作的信息时，我们需要使用grad_fn的next_functions属性。

3.2 梯度计算流程

• 前向传播记录操作
• 构建反向计算图
• 反向传播计算梯度
• 梯度累积到叶子节点

3.3 实际例子

import torch
x = torch.tensor(2.0, requires_grad=True)
# 构建计算图
fx = x**2 + 3*x + 1
fx.retain_grad() # fx是中间变量，默认不保留梯度，这里需要保留梯度，方便实验
gx = fx**3+4*fx+2
gx.retain_grad()
# 自动反向传播
gx.backward()

# 前向传播时：
print("x:{}".format(x))
print("fx:{}".format(fx))
print("gx:{}".format(gx))
# 反向传播时：
print("gx.grad:{}".format(gx.grad))
print("fx.grad:{}".format(fx.grad))
print("x.grad:{}".format(x.grad))
## 输出结果为：
# x:2.0
# fx:11.0
# gx:1377.0
# gx.grad:1.0
# fx.grad:367.0
# x.grad:2569.0       

刚刚的例子中， 我们定义了两个函数： \(f(x) = x^2 + 3x + 1\) 和 \(g(x) = f(x)^3 + 4f(x) + 2\)

当x输入2时，\(f(x) = 11\)，\(g(x) = 1377\)

现在我们需要计算 \(g(x)\) 的梯度，即 \(\frac{dg(x)}{dx}\)

根据链式法则，我们有：

\[\frac{dg(x)}{dx} = \frac{dg(x)}{df(x)} * \frac{df(x)}{dx}\]

其中： \(\frac{dg(x)}{df(x)}=3*f(x)^2+4=3*11^2+4=367\)

接下来： \(\frac{df(x)}{dx}=2x+3=2*2+3=7\) 所以： \(\frac{dg(x)}{dx}=367*7=2569\)

这个结果与我们的计算结果一致。所以，x上面的梯度，是反向传播到x的累计梯度，值为2569

3.4 运算定义

在刚刚的例子中，我们定义了通过一些简单运算，完成的一个计算图。并在这个计算图上，完成的正向传播和反向传播。 但是，在实际的深度学习中，我们通常会使用一些更复杂的运算，比如卷积、池化、激活函数等。这些运算的梯度计算，是通过一些特殊的规则来实现的。 下面，我们来介绍一些特殊的运算的梯度计算规则。

3.4.1 卷积运算

卷积运算的梯度计算，是通过卷积核的翻转来实现的。 在PyTorch中，卷积层的梯度通过反向传播自动计算。 —

权重的梯度由输入数据与输出梯度的互相关（cross-correlation）计算。 \(\frac{\partial L}{\partial W} = \text{conv}(X, \frac{\partial L}{\partial Y})\)
其中，\(X\) 是输入数据，\(\frac{\partial L}{\partial Y}\) 是损失对输出的梯度。

代码示例

以下示例展示了一个简单的2D卷积层，手动计算梯度并与PyTorch结果对比。

import torch
import torch.nn as nn

# 创建卷积层（无偏置，固定权重为2）
conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=2, bias=False)
with torch.no_grad():
    conv.weight.data = torch.rand(1, 1, 2, 2)

# 输入数据（全1的3x3矩阵）
x = torch.rand(1, 1, 3, 3)

print("输入：\n", x)
print("权重：\n", conv.weight)

# 前向传播
output = conv(x)
print("输出：\n", output)
loss = output.sum()     # 损失函数为输出所有元素的和
print("损失：\n", loss)
# 反向传播
loss.backward()

# 打印梯度
print("权重梯度：\n", conv.weight.grad)

执行上面的代码，最终的输出为：

输入：
tensor([[[[0.5739, 0.9303, 0.5375],
[0.2379, 0.0570, 0.7423],
[0.7322, 0.7695, 0.8314]]]])
权重：
Parameter containing:
tensor([[[[0.8531, 0.5101],
[0.8575, 0.3703]]]], requires_grad=True)
输出：
tensor([[[[1.1893, 1.3916],
[1.1448, 1.3950]]]], grad_fn=<ConvolutionBackward0>)
损失：
tensor(5.1207, grad_fn=<SumBackward0>)
权重梯度：
tensor([[[[1.7992, 2.2671],
[1.7966, 2.4002]]]])

下面手动计算一下梯度：

\(Output_{00} = X_{00}*W_{00} + X_{01}*W_{01} + X_{10}*W_{10} + X_{11}*W_{11}\) \(Output_{01} = X_{01}*W_{00} + X_{02}*W_{01} + X_{11}*W_{10} + X_{12}*W_{11}\) \(Output_{10} = X_{10}*W_{00} + X_{11}*W_{01} + X_{20}*W_{10} + X_{21}*W_{11}\) \(Output_{11} = X_{11}*W_{00} + X_{12}*W_{01} + X_{21}*W_{10} + X_{22}*W_{11}\) \(\frac{\partial L}{\partial W_{00}} = X_{00} + X_{01} + X_{10} + X_{11} = 0.5739 + 0.9303 + 0.2379 + 0.0570 = 1.7992\) \(\frac{\partial L}{\partial W_{01}} = X_{01} + X_{02} + X_{11} + X_{12} = 0.9303 + 0.5375 + 0.0570 + 0.7423 = 2.2671\) \(\frac{\partial L}{\partial W_{10}} = X_{10} + X_{11} + X_{20} + X_{21} = 0.2379 + 0.0570 + 0.7322 + 0.7695 = 1.7966\) \(\frac{\partial L}{\partial W_{11}} = X_{11} + X_{12} + X_{21} + X_{22} = 0.0570 + 0.7423 + 0.7695 + 0.8314 = 2.4002\)

验证结果正确