自动求导 (Autograd) – 神经网络的引擎

自动求导 (Autograd) – 神经网络的引擎

1. 什么是自动求导 (Autograd)？

• 核心功能： torch.autograd 是 PyTorch 的自动微分引擎，它为张量上的所有操作提供自动梯度计算。这是深度学习框架的核心，因为训练神经网络依赖于通过反向传播算法计算损失函数相对于模型参数的梯度。
• 基于磁带的系统 (Tape-based System)： PyTorch 使用一种称为“动态计算图”或“反向模式自动微分”的机制。可以将其想象成一个记录器（磁带），它记录了在张量上执行的所有操作，从而构建起一个计算图。当需要计算梯度时，PyTorch 会沿着这个记录好的图反向传播，计算出每个参与运算的张量的梯度。
• 动态性： 与 TensorFlow 1.x 等静态图框架不同，PyTorch 的计算图是在运行时动态构建的。这意味着你可以使用标准的 Python 控制流语句（如 if 语句、for 循环）来改变计算路径，而计算图会相应地动态调整。这使得模型构建和调试更加灵活直观。

2. requires_grad 属性

这是 Autograd 系统的核心开关。

• 对于一个张量 t，如果 t.requires_grad 设置为 True，PyTorch 就会开始跟踪在该张量上进行的所有操作。
• 任何由设置了 requires_grad=True 的张量经过一系列运算得到的输出张量，其 requires_grad 属性也会自动变为 True。
• 默认情况下，当你创建一个新的张量时，requires_grad 是 False。

  import torch
  
  # requires_grad 默认为 False
  x = torch.tensor([1.0, 2.0, 3.0])
  print(f"x: {x}, x.requires_grad: {x.requires_grad}")
  
  # 手动设置 requires_grad 为 True
  y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
  print(f"y: {y}, y.requires_grad: {y.requires_grad}")
  
  # z 的 requires_grad 会是 True，因为 y.requires_grad 是 True
  z = x + y
  print(f"z: {z}, z.requires_grad: {z.requires_grad}")
  
  # 可以通过 .requires_grad_(True) 原地修改
  x.requires_grad_(True)
  print(f"x after x.requires_grad_(True): {x.requires_grad}")
  w = x * 2
  print(f"w: {w}, w.requires_grad: {w.requires_grad}")
  

• 叶子节点 (Leaf Nodes)： 在计算图中，由用户直接创建的、requires_grad=True 的张量通常被称为叶子节点。梯度会累积到这些叶子节点的 .grad 属性中。

3. 计算图 (Computational Graph)

• 概念： 当你对设置了 requires_grad=True 的张量执行操作时，PyTorch 会在后台构建一个有向无环图 (DAG)，称为计算图。

  • 节点 (Nodes)： 图中的节点代表张量。
  • 边 (Edges)： 图中的边代表将输入张量映射到输出张量的函数（操作）。

• grad_fn 属性：

  • 如果一个张量的 requires_grad 为 True，并且它是由某个操作产生的（即它不是叶子节点），那么它会有一个 grad_fn 属性。
  • grad_fn 指向创建该张量的函数（操作），并记录了计算梯度的反向传播路径。
  • 叶子节点的 grad_fn 为 None。

  a = torch.tensor([2.0, 3.0], requires_grad=True)
  b = torch.tensor([6.0, 4.0], requires_grad=True)
  
  Q = 3*a**3 - b**2
  
  print(f"a.is_leaf: {a.is_leaf}, a.grad_fn: {a.grad_fn}") # True, None
  print(f"b.is_leaf: {b.is_leaf}, b.grad_fn: {b.grad_fn}") # True, None
  print(f"Q.is_leaf: {Q.is_leaf}, Q.grad_fn: {Q.grad_fn}") # False, <SubBackward0 object at ...>
  

  在这个例子中：

  • a 和 b 是叶子节点。
  • Q 是通过一系列操作 (**, *, -) 从 a 和 b 计算得到的。Q.grad_fn 引用了最后一步操作 (减法 SubBackward0)，这个操作知道如何将梯度反向传播给它的输入 (即 3*a**3 和 b**2)。

4. .backward() 方法 – 计算梯度

• 当计算图构建完成，并且你有一个最终的标量输出（通常是损失函数的值）时，你可以调用该标量张量的 .backward() 方法来自动计算梯度。

• .backward() 会从调用它的张量开始，沿着计算图反向传播，并根据链式法则计算所有 requires_grad=True 的叶子节点的梯度。

• 计算得到的梯度会累积到相应叶子张量的 .grad 属性中。

  # 续上例
  # 假设 Q 是一个标量 (如果 Q 不是标量，backward() 需要一个 gradient 参数)
  # 为了演示，我们先将 Q 变成一个标量，例如通过求和
  L = Q.sum() # L 是一个标量
  print(f"L: {L}, L.grad_fn: {L.grad_fn}") # L.grad_fn is <SumBackward0 object>
  
  # 现在对标量 L 调用 backward()
  L.backward()
  
  # 查看 a 和 b 的梯度
  # dL/da = d(3a^3 - b^2)/da = 9a^2
  # dL/db = d(3a^3 - b^2)/db = -2b
  print(f"Gradient of L w.r.t. a (dL/da): {a.grad}") # 预期: 9 * [2^2, 3^2] = [36, 81]
  print(f"Gradient of L w.r.t. b (dL/db): {b.grad}") # 预期: -2 * [6, 4] = [-12, -8]
  

• 关于 .backward() 的重要说明：

  • 标量输出： backward() 通常在标量输出上调用（例如损失函数）。如果张量 y 不是标量，而你想计算 dy/dx，你需要向 y.backward(gradient) 传入一个与 y 形状相同的 gradient 张量。这个 gradient 张量通常是上游传来的梯度，或者是全1的张量（如果你想计算雅可比向量积）。对于损失函数（标量），gradient 参数默认为 torch.tensor(1.0)。
  • 梯度累积： 默认情况下，梯度是累积的。这意味着如果你多次调用 backward() 而不清除梯度，新的梯度值会加到 .grad 属性的现有值上。在训练神经网络时，通常在每次迭代（参数更新）之前，需要使用 optimizer.zero_grad() 或手动将叶子节点的 .grad 属性清零 (tensor.grad.zero_())。
  • 只能对叶子节点获取 .grad： 只有 is_leaf=True 且 requires_grad=True 的张量，在调用 backward() 后，其 .grad 属性才会被填充。中间节点的梯度在反向传播过程中被计算和使用，但通常不会被存储（除非你使用 retain_grad()）。
  • 计算图的释放： 默认情况下，为了节省内存，当 backward() 执行完毕后，计算图中的非叶子节点的部分会被释放。如果你需要多次对同一个计算图进行 backward() (例如，计算不同函数的梯度，或者在某些高级应用中)，你需要在第一次调用 backward() 时设置 retain_graph=True。

  # 梯度累积示例
  x = torch.tensor([2.0], requires_grad=True)
  y1 = x * 2
  y2 = x * 3
  
  # 第一次 backward
  y1.backward() # dy1/dx = 2
  print(f"After y1.backward(), x.grad: {x.grad}") # tensor([2.])
  
  # 第二次 backward (没有清零梯度)
  y2.backward() # dy2/dx = 3
  print(f"After y2.backward() (no zero_grad), x.grad: {x.grad}") # tensor([5.]) (2+3)
  
  # 清零梯度
  x.grad.zero_()
  y2.backward()
  print(f"After zero_grad() and y2.backward(), x.grad: {x.grad}") # tensor([3.])
  
  # retain_graph 示例
  a = torch.tensor([3.0], requires_grad=True)
  b = a * 2
  c = a * 3
  
  # 如果不设置 retain_graph=True，第二次 backward 会报错
  # b.backward() # 第一次 backward
  # c.backward() # RuntimeError: Trying to backward through the graph a second time
  
  b.backward(retain_graph=True) # 保留计算图
  print(f"After b.backward(retain_graph=True), a.grad: {a.grad}") # tensor([2.])
  a.grad.zero_() # 清零梯度
  c.backward() # 现在可以对 c 进行 backward
  print(f"After c.backward(), a.grad: {a.grad}") # tensor([3.])
  

5. 停止跟踪历史记录 (Disabling Gradient Tracking)

有时候，我们不希望 PyTorch 跟踪某些操作，例如：

• 在模型评估（推理）阶段，不需要计算梯度。
• 当我们要更新模型参数时 (例如 weight = weight - learning_rate * weight.grad)，这个更新操作本身不应该被 Autograd 跟踪。

PyTorch 提供了几种方法来停止跟踪：

• torch.no_grad() 上下文管理器：

  • 这是最常用的方法。在其作用域内的所有计算都不会被跟踪，即使输入的张量 requires_grad=True。
  • 这对于模型评估非常有用，可以显著减少内存消耗并加速计算。

  x = torch.randn(3, requires_grad=True)
  y = x * 2
  print(f"y.requires_grad (inside tracking): {y.requires_grad}") # True
  
  with torch.no_grad():
      z = x * 3 # x 仍然是 requires_grad=True
      print(f"z.requires_grad (inside no_grad): {z.requires_grad}") # False
      print(f"x.requires_grad (inside no_grad): {x.requires_grad}") # True, x 本身属性不变
  # 退出 no_grad 上下文后，跟踪恢复
  w = x * 4
  print(f"w.requires_grad (outside no_grad): {w.requires_grad}") # True
  

• .detach() 方法：

  • tensor.detach() 会创建一个与原张量共享数据的新张量，但这个新张量将从当前的计算图中分离出来，其 requires_grad 属性为 False，并且它没有 grad_fn。
  • 如果你想让某个张量不再参与梯度计算，但仍希望使用它的值，可以使用 .detach()。

  a = torch.randn(3, requires_grad=True)
  b = a * 2
  print(f"b.requires_grad: {b.requires_grad}, b.grad_fn: {b.grad_fn}")
  
  c = b.detach()
  print(f"c.requires_grad: {c.requires_grad}, c.grad_fn: {c.grad_fn}") # False, None
  
  # 修改 c 不会影响 a 的梯度计算，因为 c 被分离了
  # 修改 a 或 b 仍然会影响通过 b 的梯度计算
  print(f"c shares data with b: {c.data_ptr() == b.data_ptr()}") # True
  

• 就地修改 requires_grad 属性：

  • tensor.requires_grad_(False) 可以直接修改张量的 requires_grad 属性。
  • 警告： 如果一个张量已经参与了计算图的构建（即它有一个 grad_fn），直接修改它的 requires_grad 属性可能会导致后续 backward() 出错或行为不符合预期。通常建议在张量创建时或在它参与任何需要梯度的计算之前设置好 requires_grad。

6. Autograd 与神经网络训练

Autograd 是神经网络训练的核心。一个典型的训练步骤如下：

1. 前向传播： 将输入数据喂给模型，模型计算输出。这个过程中，如果模型的参数 (权重和偏置) 设置了 requires_grad=True，Autograd 会构建计算图。

   # 假设 model 是一个 nn.Module 实例，input_data 是输入张量
   # model.parameters() 中的张量默认 requires_grad=True
   # output = model(input_data)
   

2. 计算损失： 将模型的输出与真实标签比较，计算损失值 (一个标量)。

   # loss = criterion(output, target_labels) # criterion 是损失函数
   

3. 梯度清零： 在计算新的梯度之前，清除上一次迭代累积的梯度。

   # optimizer.zero_grad() # optimizer 是一个 torch.optim.Optimizer 实例
   

4. 反向传播： 调用损失张量的 .backward() 方法，计算损失相对于所有模型参数的梯度。这些梯度会存储在参数的 .grad 属性中。

   # loss.backward()
   

5. 参数更新： 使用优化器 (如 SGD, Adam) 根据计算出的梯度来更新模型的参数。这个步骤通常在 torch.no_grad() 环境下进行，因为参数更新本身不应该被跟踪。

   # optimizer.step()
   

7. 实践：简单的线性回归示例

让我们用 Autograd 来手动实现一个简单的线性回归，以加深理解。 目标：拟合 y = w * x + b，其中 w=2, b=1。

import torch

# 0. 准备数据
X_numpy = np.array([1, 2, 3, 4], dtype=np.float32)
Y_numpy = np.array([3, 5, 7, 9], dtype=np.float32) # 真实 y = 2*x + 1

X = torch.from_numpy(X_numpy).view(-1, 1) # (4, 1)
Y = torch.from_numpy(Y_numpy).view(-1, 1) # (4, 1)

# 1. 初始化参数 (需要计算梯度的叶子节点)
w = torch.randn(1, requires_grad=True, dtype=torch.float32)
b = torch.randn(1, requires_grad=True, dtype=torch.float32)
print(f"Initial w: {w.item()}, Initial b: {b.item()}")

# 2. 定义模型 (前向传播)
def forward(x_input):
    return x_input @ w + b # 使用 @ 进行矩阵乘法，或者 x_input * w 如果 w 是标量

# 3. 定义损失函数 (均方误差)
def loss_fn(y_pred, y_true):
    return torch.mean((y_pred - y_true)**2)

# 4. 设置学习率和迭代次数
learning_rate = 0.01
n_iters = 100

# 5. 训练循环
for epoch in range(n_iters):
    # 前向传播
    y_predicted = forward(X)

    # 计算损失
    loss = loss_fn(y_predicted, Y)

    # 反向传播 (计算梯度)
    loss.backward() # 计算 dloss/dw 和 dloss/db

    # 更新参数 (在 no_grad 环境下，因为更新操作不应被跟踪)
    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad

    # 清零梯度，为下一次迭代做准备
    w.grad.zero_()
    b.grad.zero_()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{n_iters}], Loss: {loss.item():.4f}, w: {w.item():.3f}, b: {b.item():.3f}')

print(f"\nTraining finished.")
print(f"Predicted w: {w.item():.3f} (True w: 2)")
print(f"Predicted b: {b.item():.3f} (True b: 1)")

# 使用训练好的模型进行预测
test_x = torch.tensor([[5.0]])
predicted_y = forward(test_x)
print(f"Prediction for x=5: {predicted_y.item():.3f} (True y: 11)")

在这个例子中：

• w 和 b 是我们想要学习的参数，所以 requires_grad=True。
• loss.backward() 计算了损失相对于 w 和 b 的梯度，并存储在 w.grad 和 b.grad 中。
• 我们手动使用梯度下降更新了 w 和 b。
• 每次更新后，我们调用 .grad.zero_() 来清除梯度。