2. 案例2

任务说明

本案例实现一个3×3、stride=1、same padding的二维卷积算子，包含两种路径：直接卷积（四重循环）与im2col+GEMM路径。im2col将输入特征图展开为二维矩阵，再调用矩阵乘完成卷积计算。最终对比两种路径的输出一致性，并给出可复现的性能测量方法。

实现代码

import numpy as np
import time

def direct_conv2d(inputs, weights, bias=None, stride=1, pad=1):
    """
    直接卷积实现
    inputs: (N, C, H, W)
    weights: (M, C, KH, KW)
    """
    N, C, H, W = inputs.shape
    M, _, KH, KW = weights.shape
    OH = (H + 2*pad - KH) // stride + 1
    OW = (W + 2*pad - KW) // stride + 1

    # pad input
    padded = np.pad(inputs, ((0,0),(0,0),(pad,pad),(pad,pad)), mode='constant')
    out = np.zeros((N, M, OH, OW), dtype=np.float32)

    for n in range(N):
        for m in range(M):
            for oh in range(OH):
                for ow in range(OW):
                    h_start = oh * stride
                    w_start = ow * stride
                    patch = padded[n, :, h_start:h_start+KH, w_start:w_start+KW]
                    out[n, m, oh, ow] = np.sum(patch * weights[m]) + (bias[m] if bias is not None else 0)
    return out

def im2col(inputs, KH, KW, stride=1, pad=1):
    """
    im2col展开
    inputs: (N, C, H, W)
    返回: (N * OH * OW, C*KH*KW)
    """
    N, C, H, W = inputs.shape
    OH = (H + 2*pad - KH) // stride + 1
    OW = (W + 2*pad - KW) // stride + 1
    padded = np.pad(inputs, ((0,0),(0,0),(pad,pad),(pad,pad)), mode='constant')

    cols = np.zeros((N * OH * OW, C * KH * KW), dtype=np.float32)
    idx = 0
    for n in range(N):
        for oh in range(OH):
            for ow in range(OW):
                h_start = oh * stride
                w_start = ow * stride
                patch = padded[n, :, h_start:h_start+KH, w_start:w_start+KW]
                cols[idx] = patch.flatten()
                idx += 1
    return cols

def im2col_gemm_conv2d(inputs, weights, bias=None, stride=1, pad=1):
    """
    im2col + GEMM 卷积
    """
    N, C, H, W = inputs.shape
    M, _, KH, KW = weights.shape
    OH = (H + 2*pad - KH) // stride + 1
    OW = (W + 2*pad - KW) // stride + 1

    # im2col
    cols = im2col(inputs, KH, KW, stride, pad)  # (N*OH*OW, C*KH*KW)

    # 权重矩阵: (M, C*KH*KW)
    W_matrix = weights.reshape(M, -1)

    # GEMM: (N*OH*OW, M) = cols @ W_matrix.T
    gemm_out = cols @ W_matrix.T  # 使用numpy的BLAS后端

    if bias is not None:
        gemm_out += bias

    # reshape回(N, M, OH, OW)
    out = gemm_out.reshape(N, OH, OW, M).transpose(0, 3, 1, 2)
    return out

# ============ 测试与验证 ============
if __name__ == "__main__":
    np.random.seed(42)
    N, C, H, W = 2, 3, 8, 8
    M, KH, KW = 4, 3, 3

    inputs = np.random.randn(N, C, H, W).astype(np.float32)
    weights = np.random.randn(M, C, KH, KW).astype(np.float32)
    bias = np.random.randn(M).astype(np.float32)

    # 1. 正确性验证
    out_direct = direct_conv2d(inputs, weights, bias)
    out_gemm = im2col_gemm_conv2d(inputs, weights, bias)

    diff = np.max(np.abs(out_direct - out_gemm))
    print(f"最大误差: {diff:.2e}")
    assert diff < 1e-5, "两种实现结果不一致！"
    print("输出一致性验证通过")

    # 2. 性能对比（仅演示测量方法）
    # 使用更大尺寸以获得可观测差异
    N_big, C_big, H_big, W_big = 4, 16, 32, 32
    M_big = 32
    inputs_big = np.random.randn(N_big, C_big, H_big, W_big).astype(np.float32)
    weights_big = np.random.randn(M_big, C_big, KH, KW).astype(np.float32)
    bias_big = np.random.randn(M_big).astype(np.float32)

    # 预热
    _ = direct_conv2d(inputs_big, weights_big, bias_big)
    _ = im2col_gemm_conv2d(inputs_big, weights_big, bias_big)

    # 直接卷积计时
    t0 = time.perf_counter()
    out_direct_big = direct_conv2d(inputs_big, weights_big, bias_big)
    t1 = time.perf_counter()
    direct_time = t1 - t0

    # im2col+GEMM计时
    t0 = time.perf_counter()
    out_gemm_big = im2col_gemm_conv2d(inputs_big, weights_big, bias_big)
    t1 = time.perf_counter()
    gemm_time = t1 - t0

    # 验证大尺寸下结果一致
    diff_big = np.max(np.abs(out_direct_big - out_gemm_big))
    print(f"大尺寸最大误差: {diff_big:.2e}")

    print(f"直接卷积耗时: {direct_time*1000:.2f} ms")
    print(f"im2col+GEMM耗时: {gemm_time*1000:.2f} ms")
    print(f"加速比: {direct_time/gemm_time:.2f}x")

运行与验证

直接运行脚本即可完成验证：

python conv_im2col_gemm.py

预期输出：

最大误差: 0.00e+00
输出一致性验证通过
大尺寸最大误差: 0.00e+00
直接卷积耗时: xxx ms
im2col+GEMM耗时: xxx ms
加速比: x.xx x

其中耗时与加速比取决于运行机器的CPU与BLAS库配置，数值会有所不同。判定标准：最大误差小于1e-5，且im2col+GEMM路径在较大输入上通常快于直接卷积（加速比>1）。若误差过大，请检查数据类型是否一致（需均为float32）。

关键说明

• im2col展开使用显式循环实现，便于理解数据布局；实际工程中可用np.lib.stride_tricks.as_strided或框架内置函数优化。
• GEMM调用使用numpy的@运算符，其底层自动调用BLAS库（如OpenBLAS/MKL），即为本案例所指的GEMM库调用。
• 性能对比仅演示测量方法，实际加速比受输入尺寸、通道数、卷积核大小及BLAS线程数影响。小尺寸输入下im2col因内存展开开销可能更慢，属正常现象。