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

1次阅读

没有评论

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

BP 算法通过链式法则实现了误差信号的层层反向传播，让深层神经网络的参数更新成为可能。相比于早期感知机的局限性，BP 算法能自动学习特征层级表示。没有 BP 算法，就没有现代深度学习的繁荣。

在标准 BP 实现中，前向传播的所有中间结果都需要保存以供反向计算。例如计算图 X -> Linear1 -> ReLU -> Linear2 -> Loss 需要存储：

每一层的输入输出张量
激活函数的中间状态
权重矩阵的临时计算结果

对于 batch_size=1024 的 ResNet50，中间变量可占用超过 12GB 显存。

以 ReLU 函数为例：

dReLU(x)/dx = 1 if x > 0 else 0

这意味着约 50% 的梯度计算是无效的零值，但传统实现仍然会分配内存并进行计算。

对比两种全连接层实现：

原生循环实现（FLOPs: O(n³)）

for i in range(batch_size):
    for j in range(out_features):
        for k in range(in_features):
            grad_W[j,k] += x[i,k] * grad_y[i,j]

矩阵化实现（FLOPs: O(n²)）

grad_W = x.T @ grad_y  # (in_dim, batch) @ (batch, out_dim)

实测显示矩阵运算可加速 3 - 5 倍。

PyTorch 的两种内存控制方法：

# 方法 1：释放非必要节点
with torch.no_grad():
    intermediate = layer1(x).detach()

# 方法 2：精确控制保留的图结构
loss.backward(retain_graph=True)  # 需要多次 backward 时使用

class EfficientLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(input, weight)
        return input @ weight.t()

    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_tensors
        grad_input = grad_output @ weight  # 复用权重矩阵
        grad_weight = grad_output.t() @ input
        return grad_input, grad_weight

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

total_norm = torch.norm(torch.stack([torch.norm(p.grad) for p in model.parameters()]), 
    2.0
)
clip_coef = max_norm / (total_norm + 1e-6)
if clip_coef < 1:
    for p in model.parameters():
        p.grad.detach().mul_(clip_coef)

scaler = torch.cuda.amp.GradScaler()  # 自动处理梯度缩放
with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

在 Transformer 的 self-attention 中：

$$
\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
$$

如何优化 Q /K/ V 三个大矩阵的梯度计算？是否可以利用 attention score 的稀疏性？这些思考留给读者在实践中探索。

正文完