PyTorch Autograd 原理解析：3个关键属性与动态计算图构建实战

PyTorch的自动微分机制（Autograd）是深度学习框架中最具革命性的设计之一。不同于静态计算图的框架，PyTorch采用动态计算图的方式，使得模型构建和调试过程变得异常灵活。本文将深入剖析Autograd的三大核心属性—— requires_grad 、 grad 和 grad_fn ，并通过构建可视化计算图的实战案例，帮助开发者掌握梯度计算的内在机制。

1. Autograd 核心三属性解析

1.1 requires_grad：梯度追踪开关

requires_grad 是PyTorch张量的布尔属性，决定是否对该张量进行梯度追踪。当创建一个张量时，默认情况下该属性为False：

import torch

x = torch.tensor([1.0, 2.0])
print(x.requires_grad)  # 输出: False

要启用梯度追踪，有几种常用方式：

# 创建时直接指定
x = torch.tensor([1.0, 2.0], requires_grad=True)

# 对现有张量启用
x.requires_grad_(True)

# 通过函数转换
x = torch.randn(3, requires_grad=True)

实际应用场景 ：

• 冻结预训练模型参数时，将不需要更新的参数设为 requires_grad=False
• 临时禁用梯度计算时（如模型评估阶段），使用 torch.no_grad() 上下文管理器

注意：修改 requires_grad 属性是原地操作(in-place)，会直接影响原始张量

1.2 grad：梯度值的存储容器

当执行反向传播后，计算出的梯度值会存储在张量的 grad 属性中。这个属性具有以下特点：

• 初始值为None
• 与原始张量同形状
• 梯度会累积（多次反向传播时）

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.backward(torch.tensor([1.0, 1.0]))
print(x.grad)  # 输出: tensor([2., 4.])

梯度累积示例 ：

x = torch.tensor([1.0], requires_grad=True)
for _ in range(3):
    y = x * 2
    y.backward()
print(x.grad)  # 输出: tensor([6.]) 而非 tensor([2.])

为避免梯度累积，需要在每次反向传播前手动清零：

optimizer.zero_grad()  # 优化器方式
x.grad.zero_()        # 直接操作方式

1.3 grad_fn：计算图的构建引擎

每个由操作创建的张量都会有一个 grad_fn 属性，它指向创建该张量的 Function 类实例。这些 Function 对象构成了计算图的反向传播路径。

常见操作对应的grad_fn类型：

操作类型	grad_fn 类	说明
加法	AddBackward	元素级加法
乘法	MulBackward	元素级乘法
矩阵乘	MmBackward	矩阵乘法
ReLU	ReluBackward	ReLU激活函数
求和	SumBackward	张量求和

x = torch.tensor([1.0], requires_grad=True)
y = x * 2
z = torch.relu(y)
print(z.grad_fn)          # 输出: <ReluBackward0 object>
print(z.grad_fn.next_functions)  # 输出: ((<MulBackward0 object>, 0),)

2. 动态计算图构建实战

2.1 手动构建计算图

让我们通过一个具体例子来理解动态计算图的构建过程：

# 创建叶子节点
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)

# 构建计算图
c = a * b       # c.grad_fn = <MulBackward0>
d = torch.sin(c) # d.grad_fn = <SinBackward>
e = d + 1       # e.grad_fn = <AddBackward0>

此时的计算图结构为：

a → Mul → c → Sin → d → Add → e
b ↗

2.2 可视化计算图

虽然PyTorch没有内置的可视化工具，但我们可以使用第三方库 torchviz 来生成计算图：

from torchviz import make_dot

# 构建计算图
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)
c = a * b
d = torch.sin(c)
e = d + 1

# 生成可视化图形
make_dot(e, show_attrs=True, show_saved=True)

执行这段代码会生成一个DOT格式的计算图，可以使用Graphviz工具渲染成图像。图中会清晰显示：

• 各张量的 grad_fn 类型
• 计算图的前向传播路径
• 各节点的梯度计算关系

2.3 反向传播过程解析

当调用 e.backward() 时，PyTorch会执行以下步骤：

1. 从 e.grad_fn （AddBackward）开始反向传播
2. 计算 d 的梯度：∂e/∂d = 1
3. 传递到 d.grad_fn （SinBackward），计算∂d/∂c = cos(c)
4. 传递到 c.grad_fn （MulBackward），计算∂c/∂a = b 和 ∂c/∂b = a
5. 应用链式法则，得到最终梯度：
   • ∂e/∂a = ∂e/∂d * ∂d/∂c * ∂c/∂a = 1 * cos(6) * 3
   • ∂e/∂b = ∂e/∂d * ∂d/∂c * ∂c/∂b = 1 * cos(6) * 2

3. 梯度控制技术对比

3.1 torch.no_grad() vs detach()

两种技术都用于阻止梯度计算，但实现机制不同：

特性	torch.no_grad()	tensor.detach()
作用范围	上下文管理器内的所有操作	特定张量
内存占用	不保存中间计算结果	保留原始值但不记录操作
典型用途	模型评估、推理阶段	从计算图中提取中间结果

代码示例对比 ：

# no_grad示例
with torch.no_grad():
    y = x * 2  # y.requires_grad = False

# detach示例
y = x.detach() * 2  # y.requires_grad = False

3.2 retain_graph机制

默认情况下，反向传播后计算图会被释放。如果需要多次反向传播，需设置 retain_graph=True ：

x = torch.tensor([1.0], requires_grad=True)
y = x ** 2

# 第一次反向传播
y.backward(retain_graph=True)
print(x.grad)  # tensor([2.])

# 第二次反向传播（不报错）
y.backward()
print(x.grad)  # tensor([4.]) (梯度累积)

应用场景 ：

• GAN训练中需要交替更新生成器和判别器
• 某些二阶导数计算场景
• 复杂损失函数的多阶段优化

4. 实战：自定义Autograd Function

PyTorch允许通过继承 torch.autograd.Function 创建自定义自动微分操作：

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

# 使用自定义Function
x = torch.randn(3, requires_grad=True)
y = MyReLU.apply(x)
y.backward(torch.ones_like(y))

自定义Function必须实现两个静态方法：

• forward : 定义前向计算
• backward : 定义梯度计算规则

关键点 ：

• ctx.save_for_backward 保存反向传播需要的张量
• backward 的输入是输出梯度的张量
• 返回的梯度数量应与 forward 输入参数数量一致

5. 常见问题与调试技巧

5.1 梯度消失/爆炸诊断

当遇到训练不稳定问题时，可以添加梯度监控：

# 注册钩子记录梯度
def grad_hook(grad):
    print(f"Gradient norm: {grad.norm().item():.4f}")

x = torch.randn(3, requires_grad=True)
h = x.register_hook(grad_hook)  # 每次计算x的梯度时调用

y = x ** 2
y.backward(torch.ones_like(y))

# 移除钩子
h.remove()

5.2 非标量输出的反向传播

当输出不是标量时，需要提供 gradient 参数：

x = torch.randn(3, requires_grad=True)
y = x * 2

# 正确方式1：提供gradient参数
y.backward(torch.tensor([1.0, 1.0, 1.0]))

# 正确方式2：先求和
y.sum().backward()

5.3 内存优化技巧

对于大型模型，可以使用以下技术减少内存占用：

# 使用checkpoint节省内存
from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # 复杂的计算过程
    return x ** 2

x = torch.randn(10, requires_grad=True)
y = checkpoint(custom_forward, x)  # 不保存中间结果
y.backward(torch.ones_like(y))