AI

从零开始的 CUDA 神经网络算子

从零开始的 CUDA 神经网络算子

[!CAUTION]

本笔记仅供参考，请勿抄袭。

CUDA 算子简要介绍

Lab 3 要在上一份 Tensor 上补全连接、二维卷积、最大池化、Softmax 和交叉熵。逐元素 kernel 已经不够用了：全连接与卷积需要矩阵乘，卷积还要处理窗口展开，池化反向会遇到重复写入，Softmax 则要防止指数溢出。

在动手之前

全连接层和卷积层的主要计算都能归结为矩阵乘 。cuBLAS 使用列主序解释矩阵，而 C++ 数组通常按行主序保存，调用 cublasSgemm 时要同时处理操作数顺序、转置标志和 leading dimension。shape 正确并不代表布局正确，测试时必须写出每个矩阵的逻辑维度。

卷积无法直接把 NCHW 输入交给 GEMM。im2col 会把每个感受野展开为一列，卷积核展开为行，矩阵乘完成后再把结果恢复成输出布局。反向中的 col2im 需要把多个窗口对同一输入位置的贡献累加，重叠窗口会带来写冲突。

最大池化反向要找回前向选中的位置；这份实现没有保存 argmax，而是重新扫描窗口。Softmax 先减去每行最大值以抑制指数溢出，交叉熵与 Softmax 合并后，logits 梯度可以直接写成概率与 one-hot 标签之差。

本次算子同时包含逐元素操作、归约、矩阵乘和数据重排。逐元素 kernel 使用一线程一元素，归约可以先借助 Thrust，矩阵乘交给 cuBLAS，窗口展开与偏置广播则由独立 kernel 完成。

开始动手！

1
2
3
4
5

Lab3/
├── tensor.h           Tensor 与 CUDA 工具声明
├── tensor.cu          上一份 Lab 的基础算子
├── tensor_kernel.h    本次新增的神经网络算子
└── main.cu            五组独立的小规模测试

全连接、卷积、池化、Softmax 和交叉熵分别验收。每组测试独立分配输入与梯度，某个算子失败时不会牵连其他模块。

tensor.cu 保留存储与基础 Tensor 行为，新增实现集中在 tensor_kernel.h。算子接口拿到的是裸指针和 shape 参数，输出空间仍由调用端准备。

cuBLAS 与矩阵布局

卷积和全连接看起来不同，展开后都依赖矩阵乘。代码用一层 sgem() 包装 cublasSgemm，上层继续按照行主序的 $C=AB$ 传入形状和转置标志，适配层再转换成 cuBLAS 的列主序参数。

BLAS 的三个层次

cublasSgemm 属于 BLAS Level 3，完成

$$ C\leftarrow\alpha AB+\beta C $$

全连接前向、输入梯度和权重梯度都能调用它；卷积先经过 im2col，随后也落到同一个接口。偏置广播、偏置归约和数据重排仍由小型 CUDA kernel 完成。

sgem() 接收的 m、n、k 描述逻辑上的行主序矩阵，内部用两次 cublasSgeam 转置结果，并调整 cublasSgemm 的转置标志。调用处因而可以继续写熟悉的 $C=AB$ ，代价是每次矩阵乘都会申请一块 $m\times n$ 的临时显存。

当前适配层会创建临时矩阵并做显式转置。源码中的 temp 没有在返回前 cudaFree()，每调用一次 GEMM 就会泄漏一块结果大小的显存。行主序与列主序之间还满足

$$ (AB)^\mathsf{T}=B^\mathsf{T}A^\mathsf{T} $$

利用这个关系可以交换操作数和维度，省掉显式转置。

当前每次全连接、卷积都会 cublasCreate()，结束后再销毁。这样写起来直接，但初始化开销也进入了每次调用。

全连接

输入、权重和偏置的形状为

$$ X\in\mathbb{R}^{N\times C_{\mathrm{in}}},\qquad W\in\mathbb{R}^{C_{\mathrm{out}}\times C_{\mathrm{in}}},\qquad b\in\mathbb{R}^{C_{\mathrm{out}}} $$

前向模块先用 GEMM 计算主体，再用全一向量把偏置广播到 batch

$$ Y=XW^\mathsf{T}+b $$

反向模块返回三个彼此独立的结果

$$ \mathrm{d}X=\mathrm{d}Y W $$ $$ \mathrm{d}W=\mathrm{d}Y^\mathsf{T}X $$ $$ \mathrm{d}b=\sum_{n=1}^{N}\mathrm{d}Y_n $$

源码中的四次关键调用可以直接按 shape 对照。

结果左矩阵右矩阵shape
$Y$$X$$W^\mathsf{T}$$N\times C_{\mathrm{out}}$
$\mathrm{d}X$$\mathrm{d}Y$$W$$N\times C_{\mathrm{in}}$
$\mathrm{d}W$$\mathrm{d}Y^\mathsf{T}$$X$$C_{\mathrm{out}}\times C_{\mathrm{in}}$
$\mathrm{d}b$$\mathrm{d}Y^\mathsf{T}$$\mathbf{1}$$C_{\mathrm{out}}\times1$

偏置前向与反向都借助全一向量完成。前向用 $\mathbf{1}b^\mathsf{T}$ 广播，反向用 $\mathrm{d}Y^\mathsf{T}\mathbf{1}$ 沿 batch 求和。这样不必再写两个形状相近的小 kernel。

卷积

卷积模块由形状计算、im2col、GEMM、偏置和 col2im 组成。输入采用 NCHW 布局，卷积核依次保存输出通道、输入通道、高和宽。输出尺寸先由接口层算出

$$ H_{\mathrm{out}}=\left\lfloor \frac{H+2P_H-K_H}{S_H} \right\rfloor+1 $$ $$ W_{\mathrm{out}}=\left\lfloor \frac{W+2P_W-K_W}{S_W} \right\rfloor+1 $$

展开窗口

im2col 把每个输出位置对应的感受野展成一列

$$ X_{\mathrm{col}}\in \mathbb{R}^{(CK_HK_W)\times(H_{\mathrm{out}}W_{\mathrm{out}})} $$

卷积核展平为

$$ W_{\mathrm{row}}\in \mathbb{R}^{K\times(CK_HK_W)} $$

单张图像的前向便成为

$$ Y_{\mathrm{col}}=W_{\mathrm{row}}X_{\mathrm{col}}+b $$

Im2col 的数据布局

im2col kernel 由一个线性下标反推出通道、卷积核位置和输出坐标。落在 padding 区域的元素写零，其余位置从输入取值。

1
2
3
4
5
6
7
8

int h_in = h_out * stride_h - pad_h + k_h;
int w_in = w_out * stride_w - pad_w + k_w;

if (h_in >= 0 && h_in < H && w_in >= 0 && w_in < W) {
    output[index] = input[c * H * W + h_in * W + w_in];
} else {
    output[index] = 0.0f;
}

窗口寻址都放在 im2col 中，后面的 GEMM 不再处理 stride、padding 和越界。重叠窗口会重复保存像素，因此临时矩阵可能远大于原输入。

线性下标的拆解顺序必须和列矩阵布局完全一致。源码依次取出 w_out、h_out、k_w、k_h、c，因此最后两维是连续的输出位置。若只交换 k_w 与 h_out 的顺序，矩阵大小仍然正确，GEMM 也会正常返回，但每列已经不再是一块真实的感受野。

写回梯度

输入梯度先在列空间计算

$$ \mathrm{d}X_{\mathrm{col}} =W_{\mathrm{row}}^\mathsf{T}\mathrm{d}Y_{\mathrm{col}} $$

col2im 再把各列写回输入布局。一个像素可能被多个卷积窗口覆盖，所以写回必须累加，不能把列矩阵直接 reshape。这里没有让每个列元素去 atomicAdd 输入，而是让一个线程负责一个输入元素，在该线程内部枚举所有能覆盖它的窗口。写入地址互不冲突，代价是每个线程要做两层卷积核循环。

1
2
3
4
5
6
7
8
9

if (h_out % stride_h == 0 && w_out % stride_w == 0) {
    h_out /= stride_h;
    w_out /= stride_w;

    if (0 <= h_out && h_out < out_h &&
        0 <= w_out && w_out < out_w) {
        val += col_grad[col_index];
    }
}

整除检查不能省略。只有与 stride 网格对齐的窗口才真正出现过；若直接做整数除法，未对齐的位置也会被错误地归到附近窗口。

权重和偏置由整个 batch 共用。源码按样本循环，每张图像都重新做一次 im2col。计算 $\mathrm{d}W$ 和 $\mathrm{d}b$ 时，sgem() 的 beta 取 1，把本张图的结果加到已有梯度；计算 $\mathrm{d}X_{\mathrm{col}}$ 时 beta 取 0，因为每个样本都有独立的输入梯度区域。

进入循环前必须把 weight_grad 和 bias_grad 清零，否则 beta=1 会把未初始化显存也加进去。只用 batch size 为 1 的测试很难发现累计错误，至少要补一个两张图的例子。

最大池化

池化前向寻找每个窗口的最大值，反向把上游梯度送回最大值位置。

Max Pooling 的反向传播

当前实现没有保存 argmax，而是在反向时再次比较 value == max_value。窗口中出现多个相同最大值时，每个相等位置都会收到完整梯度；这一点与只记录一个最大值下标的实现不同。

池化窗口重叠时，多个输出也可能把梯度送到同一输入位置。此处和 col2im 一样，需要累加而非覆盖。

当前 max-pool kernel 的索引只有 $C\times H_{\mathrm{out}}\times W_{\mathrm{out}}$ ，没有 batch 维，main.cu 也只测试单张输入。接进网络以前需要在 Host 端逐样本调用，或者把 $N$ 加进 kernel 下标。否则第一张图之后的数据根本不会被处理。

反向实现直接对 input_grad[...] += grad_val。当窗口互不重叠时没有冲突；stride 小于池化核时，不同线程可能同时写同一个输入位置，需要 atomicAdd，或者改成与 col2im 类似的一输入一线程写法。

Softmax 与交叉熵

Softmax 先减去每行最大值，再计算指数和归一化

$$ p_{n,c}= \frac{\exp(z_{n,c}-\max_j z_{n,j})} {\sum_j\exp(z_{n,j}-\max_k z_{n,k})} $$

交叉熵对 batch 取平均

$$ L=-\frac{1}{N}\sum_{n=1}^{N}\log p_{n,y_n} $$

两者的反向合并后为

$$ \frac{\partial L}{\partial z_{n,c}} =\frac{p_{n,c}-\mathbb{1}[c=y_n]}{N} $$

合并接口省去了 Softmax 雅可比矩阵，也避免先得到极小概率再取对数。这里仍要统一标签表示：类别下标和 one-hot Tensor 对应不同的 kernel 输入，若 Python 层和 CUDA 层各自假设一种格式，很容易读错显存。

Softmax 没有手写归约 kernel，而是对 batch 中每一行依次调用 Thrust：先 reduce 求最大值，再 transform 求指数，第二次 reduce 求和，最后归一化。实现很短，不过每一行都会产生多次独立调度。

交叉熵前向让每个样本线程计算自身损失，再用 atomicAdd 加到一个标量中。因此调用前必须把 loss 清零。若上一轮的 loss 留在显存里，本轮输出会在旧值上继续累加，单看曲线很像训练突然发散。

成品代码

最终版本保留 Lab 2 的 tensor.h 与 tensor.cu，新增算子集中在 tensor_kernel.h，五组测试写在 main.cu。完整代码见 Lab 3 源码 。

算子验证

main.cu 为五类计算准备固定随机种子和小尺寸输入，分别检查前向、输入梯度、参数梯度与输出形状。卷积、全连接和池化的上游梯度直接填 1，归约结果可以人工核对；Softmax 每行概率和应接近 1，合并后的 logits 梯度每行之和应接近 0。五组测试彼此独立，算子末尾的 cudaDeviceSynchronize() 会及时暴露越界，也会让测试中的执行完全串行。

  • 全连接层两组输出分别为 $(1.1532,0.9037)$ 和 $(1.3388,1.1262)$ ，偏置梯度为 $(2,2)$ 。
  • 卷积输出保持 $3\times4$ ，首行为 $(1.4922,1.7930,1.8034,0.8262)$ ，偏置梯度为 12。
  • 最大池化输出为 $(0.9748,0.8535;0.6530,0.6245)$ ，每个 $2\times2$ 窗口只有最大值位置接收单位梯度。
  • Softmax 每行概率和为 1，交叉熵对 logits 的梯度每行和为 0。
1
2
3

nvcc -std=c++17 --extended-lambda -Xcompiler=/utf-8 \
  main.cu tensor.cu -lcublas -lcurand -o build/lab3.exe
./build/lab3.exe