基本概念 - 自动求导 - autograd 原理与使用

PyTorch 《基本概念 - 自动求导 - autograd 原理与使用》

在深度学习领域，梯度计算是优化模型参数的核心步骤。PyTorch 作为一个强大的深度学习框架，提供了自动求导机制 autograd，极大地简化了梯度计算的过程。本文将深入探讨 PyTorch 中 autograd 的原理与使用方法。

自动求导的基本概念

在深度学习中，我们通常需要最小化一个损失函数来优化模型的参数。为了找到损失函数的最小值，我们需要计算损失函数关于模型参数的梯度。梯度表示了函数在某一点的变化率，沿着梯度的反方向更新参数可以使损失函数的值逐渐减小。

手动计算梯度是一件非常繁琐且容易出错的事情，尤其是在模型结构复杂的情况下。自动求导机制可以自动计算函数的梯度，大大提高了开发效率。

autograd 原理

计算图

autograd 的核心原理是构建计算图（Computational Graph）。计算图是一种有向无环图（DAG），它记录了所有的计算操作和数据流动。在 PyTorch 中，每个张量（Tensor）都可以看作是计算图中的一个节点，而每个操作（如加法、乘法等）则是图中的边。

当我们进行一系列的计算时，autograd 会自动构建一个计算图。例如，对于表达式 y = a * b + c，计算图会记录 a、b、c 这三个张量以及乘法和加法操作。

前向传播和反向传播

• 前向传播：在计算图中，从输入节点开始，按照操作的顺序依次计算每个节点的值，直到得到最终的输出节点。例如，对于上述表达式，前向传播的过程就是先计算 a * b，然后将结果与 c 相加得到 y。
• 反向传播：在得到最终的输出值后，我们需要计算损失函数关于每个参数的梯度。反向传播从输出节点开始，沿着计算图的反向路径，根据链式法则依次计算每个节点的梯度。链式法则是微积分中的一个重要法则，它允许我们通过组合多个函数的导数来计算复合函数的导数。

叶子节点和非叶子节点

在计算图中，叶子节点是指那些不需要通过其他节点计算得到的张量，通常是用户直接创建的张量。非叶子节点是指通过其他节点计算得到的张量。在反向传播过程中，只有叶子节点的梯度会被保留，非叶子节点的梯度会在计算完成后被释放以节省内存。

autograd 的使用

启用自动求导

在 PyTorch 中，我们可以通过设置 requires_grad=True 来启用张量的自动求导功能。例如：

import torch
# 创建一个需要求导的张量
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)
# 进行计算
y = a * b
# 计算梯度
y.backward()
# 打印梯度
print(f"Gradient of a: {a.grad}")  # 输出: Gradient of a: 3.0
print(f"Gradient of b: {b.grad}")  # 输出: Gradient of b: 2.0

在上述代码中，我们创建了两个需要求导的张量 a 和 b，然后进行了乘法运算得到 y。调用 y.backward() 方法会自动计算 y 关于 a 和 b 的梯度，并将结果存储在 a.grad 和 b.grad 中。

非标量输出的反向传播

当输出 y 是一个标量时，我们可以直接调用 y.backward() 进行反向传播。但当 y 是一个非标量张量时，我们需要传入一个与 y 形状相同的张量作为参数，这个张量通常被称为 grad_output，用于指定每个元素的权重。例如：

import torch
# 创建需要求导的张量
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2
# 定义 grad_output
grad_output = torch.tensor([1.0, 1.0, 1.0])
# 进行反向传播
y.backward(grad_output)
# 打印梯度
print(f"Gradient of x: {x.grad}")  # 输出: Gradient of x: tensor([2., 2., 2.])

禁用自动求导

在某些情况下，我们可能不需要进行自动求导，例如在推理阶段。可以使用 torch.no_grad() 上下文管理器来临时禁用自动求导功能，这样可以减少内存消耗并提高计算速度。例如：

import torch
# 创建需要求导的张量
x = torch.tensor(2.0, requires_grad=True)
# 启用自动求导
with torch.no_grad():
    y = x * 2
# 检查 y 是否需要求导
print(f"Requires grad of y: {y.requires_grad}")  # 输出: Requires grad of y: False

总结

通过 autograd 机制，PyTorch 为我们提供了一种方便、高效的方式来计算梯度。掌握 autograd 的原理和使用方法对于深度学习的开发至关重要。希望本文能帮助你更好地理解和使用 PyTorch 中的自动求导功能。