BP算法反向传播的工程实践：从数学推导到高效实现

1次阅读

共计 1758 个字符，预计需要花费 5 分钟才能阅读完成。

BP 算法是神经网络训练的基石，它通过链式法则将误差信号从输出层传递到每一层参数。这种反向传播机制使得深层网络的参数更新成为可能。没有 BP 算法，现代深度学习的发展将难以想象。

在标准 BP 实现中，前向传播产生的所有中间变量（如 ReLU 激活前的值）都需要存储在显存中：

# 计算图示例
x = input @ w1  # 存储 x 用于反向传播
h = relu(x)     # 存储 h
output = h @ w2

对于 ResNet50 这类模型，中间变量可占用超过 6GB 显存。

以 ReLU 函数为例，其梯度为：
$$\frac{dReLU(x)}{dx} = \begin{cases}
1 & x > 0\
0 & x \leq 0
\end{cases}$$
约有 50% 的神经元梯度为 0，但传统实现仍会进行完整矩阵运算。

对比两种实现方式的 FLOPs（以全连接层为例）：

原始循环实现（假设输入维度 1000，输出 500）：

# 约 1,000,000 次乘加运算
for i in range(500):
    for j in range(1000):
        grad_w[i,j] = grad_output[i] * input[j]

矩阵运算实现：

# 单次矩阵乘法替代
grad_w = grad_output.T @ input  # 等效 FLOPs 但并行度高 10 倍

PyTorch 中的两个关键操作：

# 释放不需要的中间变量
x = x.detach()  # 从计算图剥离

# 控制计算图保留
loss.backward(retain_graph=True)  # 适用于多任务学习

import torch
import torch.nn as nn

class EfficientLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # 手动实现前向传播，控制存储内容
        self.x = x.detach()  # 不保留输入的计算图
        return x @ self.weight.T + self.bias

    def backward(self, grad_output):
        # 手动实现反向传播
        grad_weight = grad_output.T @ self.x
        grad_bias = grad_output.sum(0)
        grad_input = grad_output @ self.weight
        return grad_input, grad_weight, grad_bias

def print_memory_usage():
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"Allocated: {allocated:.2f}MB, Reserved: {reserved:.2f}MB")

推荐统计历史梯度范数：

grad_norm = torch.norm(torch.stack([p.grad.norm() for p in model.parameters()]), 
    p=2
)
clip_value = 0.5 * grad_norm.item()  # 动态阈值

注意检查数值稳定性：

with torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(input)
    # 必须保持 loss 计算在 float32
    loss = loss_fn(output.float(), target)

在 Transformer 架构中，attention 层的 QKV 计算会产生 O(N^2) 的中间变量。能否通过以下方式优化：