深度学习中的affine层反向传播：原理剖析与高效实现

1次阅读

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

Affine 层（全连接层）的前向传播可表示为：

$$ y = Wx + b $$

其中 $W \in \mathbb{R}^{m \times n}$ 为权重矩阵，$x \in \mathbb{R}^n$ 为输入向量，$b \in \mathbb{R}^m$ 为偏置项。反向传播需要计算三个梯度分量：

对输入的梯度：
$$ \frac{\partial L}{\partial x} = W^T \frac{\partial L}{\partial y} $$
对权重的梯度（需累加 batch 内所有样本）：
$$ \frac{\partial L}{\partial W} = \frac{\partial L}{\partial y} x^T $$
对偏置的梯度：
$$ \frac{\partial L}{\partial b} = \sum_{batch} \frac{\partial L}{\partial y} $$

def affine_backward_naive(dout, x, W, b):
    """
    dout: 上游梯度 (N, M)
    x: 输入数据 (N, n)
    W: 权重矩阵 (n, M)
    b: 偏置项 (M,)
    """
    dx = np.zeros_like(x)
    dW = np.zeros_like(W)
    db = np.zeros_like(b)

    # 显式循环计算（效率最低）for i in range(x.shape[0]):
        dx[i] = W.T.dot(dout[i])
        dW += np.outer(dout[i], x[i])
        db += dout[i]

    return dx, dW, db

def affine_backward_vectorized(dout, x, W, b):
    dx = dout.dot(W.T)  # (N,M) x (M,n) -> (N,n)
    dW = x.T.dot(dout)   # (n,N) x (N,M) -> (n,M)
    db = np.sum(dout, axis=0)
    return dx, dW, db

性能对比（Intel i7-8700K, batch_size=128）：

实现方式	耗时(ms)
原生 Python	12.4
NumPy 向量化	0.8

NumPy 默认使用 C 顺序（row-major）内存布局，而 BLAS 库更擅长处理 F 顺序（column-major）数据。通过显式指定内存顺序可提升 15%-20% 性能：

x = np.ascontiguousarray(x, dtype=np.float32)
W = np.asfortranarray(W, dtype=np.float32)

使用 OpenBLAS 或 MKL 的线程优化：

import os
os.environ["OMP_NUM_THREADS"] = "4"  # 根据 CPU 核心数调整

优化后性能（相同硬件）：

优化手段	耗时(ms)
基础向量化	0.8
内存布局优化	0.65
BLAS 线程优化	0.52

import numpy as np
from timeit import timeit

class AffineLayer:
    def __init__(self, input_dim, output_dim):
        self.W = np.random.randn(input_dim, output_dim) * 0.01
        self.b = np.zeros(output_dim)
        # 初始化梯度缓存
        self.dW = None
        self.db = None

    def forward(self, x):
        """
        x: (N, input_dim)
        返回: (N, output_dim)
        """
        self.x = np.ascontiguousarray(x, dtype=np.float32)
        self.W = np.asfortranarray(self.W, dtype=np.float32)
        return self.x.dot(self.W) + self.b

    def backward(self, dout):
        """
        dout: (N, output_dim)
        返回: dx (N, input_dim)
        """
        dx = dout.dot(self.W.T)
        self.dW = self.x.T.dot(dout)
        self.db = np.sum(dout, axis=0)
        return dx

测试不同 batch_size 下的耗时（单位：ms）：