深度学习核心技术:PyTorch自动微分原理与实践
在深度学习领域,优化算法的核心在于对模型参数进行求导。虽然基本的微积分原理相对简单,但面对日益复杂的神经网络模型,手动计算并更新成千上万个参数的梯度无疑是繁琐且极易出错的。为了解决这一挑战,现代深度学习框架引入了自动微分(Automatic Differentiation)技术。
这项技术能够根据用户定义的计算流程,自动构建一个计算图(Computational Graph),精确追踪每个数据如何经过一系列运算生成最终结果。随后,通过反向传播(Backpropagation)机制,系统沿着计算图的逆向路径,高效地计算出每个参数相对于目标函数的偏导数,从而极大地简化了模型训练过程。
一个简单的例子
作为一个演示,假设我们希望计算函数 y = 2 * xTx 对列向量 x 的导数。首先,我们创建一个张量 x 并为其分配初始值。
import torch
# 创建一个向量张量,并启用梯度追踪
# requires_grad=True 告诉PyTorch需要为这个张量计算梯度
input_vector = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
print("初始输入向量:", input_vector)
# 定义一个标量函数:output_scalar = 2 * (input_vector^T * input_vector)
# 这等价于 output_scalar = 2 * sum(input_vector_i^2)
output_scalar = 2 * torch.dot(input_vector, input_vector)
print("计算结果 output_scalar:", output_scalar)
# 执行反向传播,计算梯度。
# 对于标量输出,直接调用 .backward() 即可。
output_scalar.backward()
# 检查计算出的梯度,它存储在张量的 .grad 属性中
print("output_scalar 对 input_vector 的梯度:", input_vector.grad)
# 我们可以手动验证梯度:函数 y = 2x^Tx 的梯度应为 4x。
# 在本例中,input_vector 为 [1.0, 2.0, 3.0, 4.0],期望梯度为 [4.0, 8.0, 12.0, 16.0]
expected_grad = 4 * input_vector
print("期望梯度:", expected_grad)
print("计算梯度与期望梯度是否一致:", torch.allclose(input_vector.grad, expected_grad))
需要注意的是,PyTorch 默认会累积梯度。这意味着如果在一个张量上多次调用 backward() 而不清除梯度,新的梯度会累加到已有的 .grad 属性上。在进行新的梯度计算前,通常需要将梯度清零。
# 清零先前的梯度,准备进行新的计算
input_vector.grad.zero_()
print("梯度清零后:", input_vector.grad)
# 另一个简单函数示例:计算张量中所有元素的和
sum_output = input_vector.sum()
print("元素求和结果 sum_output:", sum_output)
# 执行反向传播
sum_output.backward()
print("sum_output 对 input_vector 的梯度:", input_vector.grad)
# 对于 sum_output = sum(x_i),其梯度应为所有元素均为 1 的向量。
非标量变量的反向传播
当目标输出 y 是一个非标量(例如向量或矩阵)时,对 y 直接调用 .backward() 需要传入一个 gradient 参数。这个参数是一个与 y 同形状的张量,它代表了我们希望计算的"向量-雅可比乘积"(vector-Jacobian product)。
然而,在深度学习实践中,我们通常关心的是批量样本中每个元素的损失函数对参数的偏导数之和。在这种情况下,一种常见的做法是对非标量输出执行求和操作,将结果变为标量,然后再对这个标量调用 .backward()。这等价于传入一个与非标量输出形状相同的全一向量作为 gradient 参数。
# 清零梯度,准备新的计算
input_vector.grad.zero_()
# 定义一个生成非标量输出的函数:element_wise_square = input_vector * input_vector
element_wise_square = input_vector * input_vector
print("element_wise_square (非标量输出):", element_wise_square)
# 方式一:对非标量输出求和,然后对和进行反向传播
# 这等价于计算所有元素的偏导数之和
element_wise_square.sum().backward()
print("通过 element_wise_square.sum().backward() 计算的梯度:", input_vector.grad)
# 验证梯度:如果 f(x) = x^2 (逐元素平方),则 df/dx = 2x。
# 对于 input_vector = [1.0, 2.0, 3.0, 4.0],期望梯度为 [2.0, 4.0, 6.0, 8.0]。
expected_grad_square = 2 * input_vector
print("期望梯度 (2x):", expected_grad_square)
print("计算梯度与期望梯度是否一致:", torch.allclose(input_vector.grad, expected_grad_square))
# 方式二:直接对非标量输出调用 backward(),并传入一个全一向量作为梯度参数
input_vector.grad.zero_() # 再次清零
element_wise_square.backward(torch.ones_like(element_wise_square))
print("通过 element_wise_square.backward(torch.ones_like(...)) 计算的梯度:", input_vector.grad)
分离计算
有时,我们希望将计算图中的某些部分"冻结"或"分离",使其不再参与梯度的反向传播。例如,当 y 是 x 的函数,而 z 又是 y 和 x 的函数时。如果我们希望计算 z 关于 x 的梯度,但将 y 视为一个常数(即不追溯 y 是如何由 x 计算得来),可以使用 .detach() 方法。
.detach() 会返回一个与原张量具有相同值的新张量,但这个新张量将不再连接到原张量的计算图中。因此,从分离后的张量开始的任何反向传播都不会回溯到其原始来源。
# 再次清零梯度
input_vector.grad.zero_()
# 原始计算:mid_result = input_vector * input_vector
mid_result = input_vector * input_vector
print("中间结果 mid_result:", mid_result)
# 分离计算图:将 mid_result 的值复制到一个新的张量 detached_val,
# 但切断 detached_val 与 mid_result 及其之前计算图的连接。
# 这意味着从 detached_val 开始的任何反向传播都不会回溯到 input_vector。
detached_val = mid_result.detach()
print("分离后的值 detached_val:", detached_val)
# 定义最终输出:final_output = detached_val * input_vector
# 此时,detached_val 被视为常数,只有 input_vector 对 final_output 贡献梯度。
final_output = detached_val * input_vector
print("最终输出 final_output:", final_output)
# 对 final_output 求和并进行反向传播
final_output.sum().backward()
# 此时,我们计算的是 d(sum(detached_val * input_vector)) / d(input_vector)。
# 由于 detached_val 被视为常数,梯度应该就是 detached_val 本身。
print("对 final_output 进行反向传播后 input_vector 的梯度:", input_vector.grad)
print("验证梯度是否等于 detached_val:", torch.allclose(input_vector.grad, detached_val))
# 如果不使用 detach(),原始的 mid_result 会继续追踪梯度。
# 现在我们可以单独对 mid_result 进行反向传播,获取它关于 input_vector 的梯度。
input_vector.grad.zero_() # 清零以计算 mid_result 的梯度
mid_result.sum().backward()
print("单独对 mid_result 进行反向传播后 input_vector 的梯度:", input_vector.grad)
# 此时,我们计算的是 d(sum(input_vector * input_vector)) / d(input_vector) = 2 * input_vector。
print("验证梯度是否等于 2 * input_vector:", torch.allclose(input_vector.grad, 2 * input_vector))
Python控制流的梯度计算
自动微分的强大之处在于,即使构建函数的计算图需要通过复杂的Python控制流(例如 if/else 条件语句、while 或 for 循环,或任意函数调用),我们仍然可以准确地计算出最终变量的梯度。这是因为深度学习框架会在运行时动态地构建计算图,记录实际执行的每一步操作。
考虑以下包含控制流的函数:
def conditional_function(input_param):
"""
一个包含Python控制流的示例函数。
计算过程中的循环次数和条件分支取决于输入值。
"""
intermediate_val = input_param * 2
# while循环,其迭代次数取决于 intermediate_val 的范数
while torch.linalg.norm(intermediate_val) < 1000:
intermediate_val = intermediate_val * 2
# if-else条件分支,最终结果取决于 intermediate_val 的和
if intermediate_val.sum() > 0:
final_result = intermediate_val
else:
final_result = 0.01 * intermediate_val # 改变系数,以略微修改函数逻辑
return final_result
# 创建一个单元素张量作为输入,并启用梯度追踪
input_scalar = torch.randn(size=(), requires_grad=True) # 使用随机值
print("初始输入 scalar:", input_scalar)
# 调用包含控制流的函数
output_val_cf = conditional_function(input_scalar)
print("函数计算结果:", output_val_cf)
# 执行反向传播
output_val_cf.backward()
# 验证梯度。请注意,这个函数在它的输入 input_scalar 中是分段线性的。
# 换句话说,对于任何 input_scalar,存在某个常量标量 k,使得 f(input_scalar) = k * input_scalar。
# 因此,df/d(input_scalar) 应该等于 f(input_scalar) / input_scalar。
print("计算出的梯度:", input_scalar.grad)
print("验证梯度是否等于 output_val_cf / input_scalar:", torch.allclose(input_scalar.grad, output_val_cf / input_scalar))
此示例证明了自动微分系统能够处理复杂的、动态变化的计算路径,大大提升了深度学习模型开发的灵活性和效率。
小结
自动微分是现代深度学习框架的核心能力,它极大地简化了梯度计算过程。其基本流程是:
- 标记变量: 指定需要计算梯度的张量(通常是模型参数),通过设置
requires_grad=True来告知系统追踪其计算历史。 - 构建计算图: 框架在执行前向传播(即模型预测)时,会隐式地构建一个有向无环图,记录所有张量及其之间的操作。
- 反向传播: 当调用目标张量(如损失函数)的
.backward()方法时,系统会遍历计算图,从输出端逆向回溯到输入端,自动计算并累积每个相关张量的梯度。 - 获取梯度: 计算出的梯度存储在相应张量的
.grad属性中,可用于优化器更新参数。
这一机制使得开发者无需手动推导复杂函数的导数,从而能够更专注于模型设计与实验。