自动微分(Autograd)#

自动微分是 PyTorch 的核心特性之一,它能够自动计算梯度,这是训练神经网络的基础。理解自动微分机制对于掌握 PyTorch 至关重要。

在深度学习中,自动求导主要用于两个方面:

  1. 在训练神经网络时计算梯度
  2. 进行反向传播算法的实现

自动微分基于链式法则(Chain Rule),复合函数 $y = f(g(x))$ 的导数为:

$$ \frac{dy}{dx} = \frac{dy}{dg} \cdot \frac{dg}{dx} $$

在深度学习中,模型由多层复合而成,自动求导沿计算图从输出到输入逐层应用链式法则,高效计算梯度。

核心机制#

  1. 前向传播:执行操作时构建有向无环图(DAG),叶子节点为输入张量,根节点为输出
  2. 记录操作:每个张量的 grad_fn 指向创建它时对应的运算信息,形成计算历史链
  3. 反向传播:调用 backward() 时,从根到叶遍历图,按链式法则累积梯度,记入.grad
  4. 动态图:每次迭代重新构建有向无环图,支持 Python 控制流(if/for)

基本使用#

启用梯度追踪#

import torch

# 创建需要梯度的张量
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
print(f"requires_grad: {x.requires_grad}")

# 进行计算
y = x ** 2
z = y.sum()

# 反向传播
z.backward()

# 查看梯度
print(f"x.grad: {x.grad}")  # [2., 4., 6.]

requires_grad 参数#

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

# 后续修改
x.requires_grad_(True)
x.requires_grad_(False)

# 使用 detach() 分离
y = x.detach()  # y 不追踪梯度

计算图#

PyTorch 使用动态计算图来追踪操作:

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

# 查看计算图
print(f"x.is_leaf: {x.is_leaf}")      # True(叶子节点)
print(f"y.is_leaf: {y.is_leaf}")      # False
print(f"z.is_leaf: {z.is_leaf}")      # False

# 查看梯度函数
print(f"y.grad_fn: {y.grad_fn}")      # <PowBackward0>
print(f"z.grad_fn: {z.grad_fn}")      # <SumBackward0>

backward() 方法#

标量输出#

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

# 对于标量,直接调用 backward()
y.backward()
print(x.grad)  # [2., 4.]

向量输出#

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2  # y 是向量

# 需要提供梯度权重(通常为全1)
y.backward(torch.ones_like(y))
print(x.grad)  # [2., 4.]

# 或者使用 sum() 先转换为标量
x.grad = None  # 清零梯度
y = x ** 2
y.sum().backward()
print(x.grad)  # [2., 4.]

梯度累积#

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

# 第一次反向传播
y1 = x ** 2
y1.backward()
print(x.grad)  # [2.]

# 第二次反向传播(梯度会累积)
y2 = x ** 3
y2.backward()
print(x.grad)  # [2. + 3. = 5.]

# 清零梯度
x.grad.zero_()

常见操作示例#

线性函数#

# y = ax + b
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x = torch.tensor(3.0)

y = a * x + b
y.backward()

print(f"dy/da: {a.grad}")  # 3.0 (x的值)
print(f"dy/db: {b.grad}")  # 1.0

多项式#

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 3 + 2 * x ** 2 + x
loss = y.sum()
loss.backward()

print(x.grad)  # [3*1² + 4*1 + 1, 3*2² + 4*2 + 1] = [8., 21.]

链式法则#

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
z = y ** 3
loss = z.sum()
loss.backward()

# dz/dx = dz/dy * dy/dx = 3y² * 2x = 6x⁵
print(x.grad)  # [6., 192.]

禁用梯度追踪#

torch.no_grad()#

x = torch.tensor([1.0, 2.0], requires_grad=True)

# 在 no_grad 上下文中,不追踪梯度
with torch.no_grad():
    y = x ** 2
    print(y.requires_grad)  # False

推理时禁用梯度#

model = ...  # 你的模型
model.eval()  # 设置为评估模式

with torch.no_grad():
    predictions = model(input_data)

detach()#

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

# detach 创建不追踪梯度的新张量
z = y.detach()
z.requires_grad  # False

# 但 y 仍然追踪梯度
y.backward()
print(x.grad)  # [2.]

高阶梯度#

PyTorch 支持高阶梯度计算:

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

# 一阶梯度
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
print(dy_dx)  # 3x² = 3

# 二阶梯度
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]
print(d2y_dx2)  # 6x = 6

自定义反向传播#

自定义函数#

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input ** 2
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * 2 * input

# 使用
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = MyFunction.apply(x)
y.sum().backward()
print(x.grad)  # [2., 4.]

梯度检查#

数值梯度验证#

def numerical_gradient(f, x, eps=1e-5):
    """计算数值梯度"""
    grad = torch.zeros_like(x)
    for i in range(x.numel()):
        x_plus = x.clone()
        x_plus.flat[i] += eps
        x_minus = x.clone()
        x_minus.flat[i] -= eps
        grad.flat[i] = (f(x_plus) - f(x_minus)) / (2 * eps)
    return grad

# 测试
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = (x ** 2).sum()

# 自动梯度
y.backward()
auto_grad = x.grad.clone()

# 数值梯度
x.grad = None
num_grad = numerical_gradient(lambda x: (x ** 2).sum(), x)

print(f"自动梯度: {auto_grad}")
print(f"数值梯度: {num_grad}")
print(f"差异: {(auto_grad - num_grad).abs().max()}")

常见问题#

1. 梯度为 None#

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

# 如果 y 不是标量,需要指定 grad_output
# y.backward()  # 错误!
y.sum().backward()  # 正确
# 或
y.backward(torch.ones_like(y))  # 正确

2. 只保留叶子节点的梯度#

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

print(x.grad)  # 有梯度
# print(y.grad)  # None(中间节点默认不保留)

3. 多次反向传播#

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

y.backward(retain_graph=True)  # 保留计算图
print(x.grad)  # [2.]

y.backward()  # 可以再次反向传播
print(x.grad)  # [4.](累积)

4. 不可微函数的梯度#

ReLU、sqrt 等在部分点不可微。PyTorch 采用约定:可微则用导数;凸函数用最小范数次梯度;未定义处可能返回 NaN。详见 Autograd 机制

性能优化#

1. 使用 detach() 避免不必要的计算#

# 不好的做法
for i in range(100):
    y = model(x)
    loss = criterion(y, target)
    loss.backward()

# 好的做法(如果不需要中间梯度)
for i in range(100):
    with torch.no_grad():
        y = model(x)
    loss = criterion(y, target)
    loss.backward()

2. 及时清零梯度#

optimizer.zero_grad()  # 在每次反向传播前清零
loss.backward()
optimizer.step()

练习#

  1. 实现函数 f(x) = x³ + 2x² + x,计算在 x=1 处的梯度
  2. 实现链式法则:y = sin(x²),计算 dy/dx
  3. 实现多变量函数:f(x, y) = x²y + xy²,计算所有偏导数
  4. 验证自动梯度与数值梯度的一致性