GPU 开销：为什么我们的 MNIST 模型在 CPU 上训练更快

GPU 可以并行处理大量计算，从而加快神经网络训练。但小模型的计算量可能不足以抵消操作调度、数据传输和同步带来的开销。

我们在 Quadro RTX 5000 GPU 和 Intel Core i9-10885H CPU 上比较了这个 MNIST 模型。在一次 Keras 测试中，CPU 更快：每轮 4.4 秒，GPU 为 6.5 秒。另一组测试中，JAX 在 CPU 上的训练循环最低耗时为 0.14 秒，但其计时范围与 Keras 不同。下面我们分析开销的来源，以及这些测量结果能说明什么。

我们先来看训练过程中究竟发生了什么。

拆解一个训练步

一个训练步包括前向传播、计算损失梯度、反向传播和更新权重。下面的代码直接计算 softmax 与交叉熵组合后的梯度，没有计算标量损失值：

# 前向传播
z1 = xb @ w1 + b1                                # 矩阵乘法 + 偏置（隐藏层）
a1 = np.maximum(0, z1)                            # ReLU 激活
z2 = a1 @ w2 + b2                                # 矩阵乘法 + 偏置（输出层）
exp_z = np.exp(z2 - z2.max(axis=1, keepdims=True))
probs = exp_z / exp_z.sum(axis=1, keepdims=True)  # softmax → 概率

# 损失的梯度
dz2 = probs.copy()
dz2[np.arange(bs), yb] -= 1                       # 预测偏离了多少
dz2 /= bs

# 反向传播——计算梯度
dw2 = a1.T @ dz2                                  # W2 的梯度
da1 = dz2 @ w2.T                                  # 向后流动的梯度
dz1 = da1 * (z1 > 0)                              # ReLU 的梯度
dw1 = xb.T @ dz1                                  # W1 的梯度

# 更新权重
w1 -= lr * dw1
b1 -= lr * dz1.sum(axis=0)
w2 -= lr * dw2
b2 -= lr * dz2.sum(axis=0)

时间都花到哪儿去了？为了搞清楚，我们用 time.perf_counter() 把每个操作包了起来：

t = time.perf_counter()
z1 = xb @ w1 + b1
timings["fwd: X @ W1 + b1"] += time.perf_counter() - t

t = time.perf_counter()
a1 = np.maximum(0, z1)
timings["fwd: ReLU"] += time.perf_counter() - t
# ...其余每个操作依此类推

GPU 操作是异步执行的，因此我们用 tf.test.experimental.sync_devices() 等待操作完成。不过，每个操作之后都同步也会增加开销，并阻止操作重叠执行。因此，这些耗时不能代表经过优化的训练循环。

以下是单个训练步的结果（batch_size=32，对 1000 步取平均，时间单位为微秒 μs——百万分之一秒，越小越好）：

在这个逐操作计时的测试中，CPU 更快。NumPy 的矩阵乘法调用了经过优化的 BLAS 例程，而 TensorFlow 的 GPU 计时还包含调度和同步。SIMD 指令让 CPU 能用一条指令处理多个数值，融合乘加指令则将乘法和加法合并。这里比较的是具体实现与计时方法的综合结果，不能据此单独判断两种设备的算术性能。

我们来看 xb @ w1 + b1 和 a1 @ w2 + b2——两层各自的矩阵乘法——以及 GPU 为何被设计成能把它们做得很快。下面是我们 NumPy 实现中的前向传播：

class HiddenLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # 矩阵乘法 + 偏置
        self.out = np.maximum(0, self.z)     # ReLU 激活
        return self.out

class OutputLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # 矩阵乘法 + 偏置
        exp = np.exp(self.z - np.max(self.z))
        self.probs = exp / np.sum(exp)       # softmax → 概率
        return self.probs

对单张图像，self.W @ x 将一个 128×784 的权重矩阵与包含 784 个像素值的向量相乘。每个神经元计算一个点积，需要 784 次乘法和 783 次加法。加上偏置后得到 128 个激活前的值，再经过 ReLU 得到隐藏层激活值：

W (128 × 784) @ x (784,) + b (128,) → z (128,)

对包含 32 张图像的批次，X 的形状是 (32, 784)。将逐张处理图像的实现中保存的权重转置，得到形状为 (784, 128) 的 W1。乘积的形状是 (32, 128)；随后再加上偏置并应用 ReLU：

      X (32 × 784)                W1 (784 × 128)            result (32 × 128)
 ┌─────────────────┐         ┌──────────────────┐       ┌──────────────────┐
 │ img1:  p1 … p784│         │  n1   n2  … n128 │       │ img1: z1  … z128 │
 │ img2:  p1 … p784│    @    │  w    w   …  w   │   =   │ img2: z1  … z128 │
 │ ...             │         │  ...  ...    ... │       │ ...              │
 │ img32: p1 … p784│         │  w    w   …  w   │       │ img32:z1  … z128 │
 └─────────────────┘         └──────────────────┘       └──────────────────┘

批处理让 32 张图像共用同一组权重，并将计算表示为一次矩阵乘法。这样可以减少单独调度的次数，也让实现有更多机会复用数据、并行计算。

这 4,096 个输出值可以独立计算，但经过优化的 GPU 矩阵乘法通常会把矩阵分块，分配给不同的线程组。不能仅凭输出值的数量推断有多少核心或流多处理器（SM）处于活跃状态。

作为示例，假设有 3,072 个 FP32 核心，频率为 1.8 GHz，每次乘加算作两次浮点运算，则理论峰值约为 11 TFLOPS。实际吞吐量取决于 GPU 型号、时钟频率和工作负载。

将乘法和加法分别计数，这次矩阵乘法约需 $2 \times 32 \times 784 \times 128 = 6{,}422{,}528$ 次浮点运算。按 11 TFLOPS 计算，仅算术部分的耗时下限约为 0.58 μs。这个估算假设能达到峰值吞吐量，且不计内存访问与调度。它不是实测内核耗时；运算次数除以峰值 FLOPS 得到的是时间，不是 GPU 利用率。

时间实际上花在哪里

GPU 训练步除了算术计算，还有以下几类开销。它们各自占多大比重，取决于工作负载和执行方式：

1. 内核调度。CPU 和运行时需要向 GPU 提交任务。对小规模操作，调度可能占用相当一部分时间。一个框架操作可能调用多个内核，编译器也可能把多个操作融合成一个内核；Python 表达式的数量并不决定内核启动次数。

2. 数据传输与同步。输入数据可能需要从 CPU 内存传到 GPU 显存。一个包含 32 张图像、每张 784 个 float32 像素的批次约占 100 KB。传输耗时取决于连接方式、内存分配和同步。在通常的 GPU 训练中，权重和梯度可以一直保留在 GPU 上，梯度不需要每步都传回 CPU。

3. Python 与框架开销。调度、内存分配、指标计算和数据处理都会耗时。在即时执行模式下，Python 参与各个操作的执行；图执行和编译可以减少大量重复工作。XLA 编译发生在首次需要编译函数或重新追踪时，并不是每一步的每个操作都要重新编译。

性能分析可以帮助区分这些开销，但必须分清主机端 API 调用耗时、GPU 执行时间和总耗时。

剖析 CUDA 开销

逐操作测量表明，开销会影响这个工作负载的表现。要进一步分析，还需要查看完整训练流程的性能记录。

为了找出答案，我们用了 NVIDIA 的 nsys（Nsight Systems）性能剖析器，它会拦截每一次 CUDA API 调用：

nsys profile -o keras-gpu-profile python mnist-keras.py
nsys stats --force-export=true keras-gpu-profile.nsys-rep

CUDA 是 NVIDIA 在你的代码与 GPU 硬件之间的软件层。当 TensorFlow 想做矩阵乘法时，它并不直接跟 GPU 对话——它调用「分配显存」「拷贝这些数据」「启动这个内核」之类的 CUDA 函数。每一次这样的调用都要经过驱动，各自都有开销。nsys 剖析器会把它们逐一记录下来，于是我们就能看清时间的去向。

记录中的 GPU 训练一轮耗时 6.5 秒，对应的 CUDA API 汇总如下。48,000 张训练图像、批大小为 32 时，每轮有 1,500 个训练步，但表中记录了 1,875 次图启动。仅凭调用次数，无法确定哪些调用属于训练、验证或初始化。

这些是主机端 CUDA API 调用的耗时，不是 GPU 内核的执行时间。调用可能与 GPU 工作或其他 CPU 线程的活动重叠。因此，不能从 6.5 秒中减去这些耗时，再把剩余时间归为 Python 开销。要做这种拆分，需要明确测量范围并分析时间线。Nsight Systems 区分了 API、队列和内核时间。

cuGraphLaunch 条目说明本次运行使用了 CUDA Graphs 来提交已记录的任务。图回放可以减少调度开销，但调用次数不足以证明每个训练批次对应一次图启动，也不能证明 GPU 计算耗时可以忽略。

在前面的 Keras 比较中，CPU 每轮耗时 4.4 秒，GPU 为 6.5 秒。在 CPU 上执行可以避免 CUDA 调用和向独立 GPU 传输数据，但仍有框架、调度和内存访问开销。

另一个 TensorFlow 微基准测试对 (32, 784) @ (784, 128) 计时，没有在每次操作之后都同步，得到的平均耗时如下：

在这个测试中，GPU 约快 2.5 倍。两个数值都包含软件和内存访问开销，都不是纯算术耗时。它们的计时方法也与第一张表不同，因此结果反转并不矛盾。

配套的编译训练步脚本在两种设备上执行的工作不同：GPU 路径会更新权重，CPU 路径只计算梯度。因此，这些训练步耗时不能用于直接比较 CPU 和 GPU。公平的比较需要执行相同的计算，并等待两种设备完成工作。

即使某个工作负载在 CPU 上运行更快，训练时 nvidia-smi 也可能显示 90–100% 的 GPU 利用率。

nvidia-smi 的 GPU 利用率表示采样区间内至少有一个内核正在执行的时间占比。它不表示被占用的核心比例，也不表示达到了峰值 FLOPS 的多少。因此，小内核连续运行时，即使没有充分发挥 GPU 的吞吐能力，利用率读数也可能很高。

让单次训练跑得更快

我们尝试了五种调整训练工作负载的方法：增大批次、在 CPU 上运行 Keras、使用 NumPy 循环、使用 JAX 编译，以及直接移植到 CuPy。它们影响的是不同部分的执行成本。

1. 加大批大小

更大的批次可以减少每轮的更新次数。48,000 张训练图像、批大小为 32 时，每轮有 1,500 步。批大小为 4,096 时，如果保留最后一个不完整批次，则有 12 步。本次测试中的 NumPy 和 JAX 循环会丢弃余下的样本，只处理 11 个完整批次，即 45,056 张图像。

更大的批次也增加了每次矩阵乘法的计算量，有助于利用 GPU 的并行能力。活跃 SM 的数量取决于选用的内核，需要通过测量确定，不能直接从批大小推算。

不过，批大小不能不加分辨地往上加——它直接影响模型准确率。更大的批带来更平滑但更不频繁的梯度更新，可能导致泛化变差。合适的批大小需要你针对自己的模型试出来。

2. Keras 跑在 CPU 上——跳过 CUDA 开销

在设备初始化之前调用 tf.config.set_visible_devices([], 'GPU')，可以让 TensorFlow 在 CPU 上运行。这样会移除执行路径中的 CUDA 工作，但不意味着总耗时会恰好减少各项 CUDA API 耗时之和。

3. 纯 NumPy——把框架开销也一并跳过

NumPy 训练循环不经过 TensorFlow 运行时，并使用 BLAS 完成矩阵乘法。但 Python 调用、数组分配和内存访问仍然有开销。

4. JAX 的 JIT——把整个步骤编译起来

JAX 的 jit会编译函数，使后续调用不必再通过 Python 逐个调度操作。编译可以融合操作、减少开销，但一个编译后的函数不一定只对应一个 GPU 内核。计时时应将编译时间与稳定运行后的耗时分开。

5. CuPy——要是干脆把 NumPy 搬到 GPU 上呢？

我们也把逐张处理图像的 NumPy 循环移植到了 CuPy。这样一来，小规模向量操作需要反复向 GPU 提交任务。这个版本训练五轮耗时 443 秒，超过对应 NumPy 循环耗时的六倍。这个结果只适用于当前实现；采用批处理的 CuPy 实现需要另行比较。

把这些合起来看

下表记录了另一组不同批大小的测试。各实现的工作负载并不相同：Keras 调用包含验证和指标计算，而 NumPy 与 JAX 循环只计训练耗时，并丢弃不完整批次。因此，这些数值反映的是现有脚本的运行时间，不能作为严格控制条件下的框架速度排名。

这里批大小为 32 的 CPU 耗时是 11.31 秒，而前面的测试为 4.4 秒。不能把不同测试的结果拼在一起计算加速比。严格的比较应保持数据、更新次数、验证工作、预热方式和计时范围一致。

Keras GPU 在这组测试中接近 2.3 秒。仅凭这张表，无法确定哪一部分限制了进一步提速。

Keras CPU 的耗时随批大小增加，从 11.31 秒降至 1.39 秒。这里计量的是测试中完整的 fit() 调用。

NumPy 在批大小为 2,048 时耗时 0.84 秒，在 4,096 时为 1.03 秒。仅凭耗时不能判断变慢的原因。

JAX 的纯训练循环在批大小为 4,096 时记录到 0.14 秒。6.26 秒除以 0.14 秒约等于 45，但这个比值混合了不同的批大小和工作负载。它不能证明在训练工作相同或验证准确率相同时获得了 45 倍加速。

对这个小模型，可以先测试减少重复调度和调整批大小的效果，再考虑更换硬件。除了每轮耗时，也应比较达到相同验证效果所需的时间。

用并行化做超参数搜索

上面这些办法让单次训练更快。但当你做超参数搜索时——试 5 个学习率，或 4 种架构——每次运行都是完全独立的。它们不共享权重、梯度或状态。

JAX 的 vmap为更新函数增加一个模型维度，让一次调用能用同一批图像更新五组独立参数。与 jit 配合后，可以避免用 Python 循环逐个处理模型。编译器决定如何执行这些操作，并不保证只使用一个内核，也不保证所有模型同时执行。

在这个例子中，所有模型的参数形状相同。下面展示向量化更新的代码，假设 x_batch 和 y_batch 已经准备好：

import jax
import jax.numpy as jnp
from jax import vmap, jit, grad, random

def init_params(key):
    # 与之前相同的模型：784→128→10，随机权重
    k1, k2 = random.split(key)
    w1 = random.normal(k1, (784, 128)) * jnp.sqrt(2.0 / 784)
    b1 = jnp.zeros(128)
    w2 = random.normal(k2, (128, 10)) * jnp.sqrt(2.0 / 128)
    b2 = jnp.zeros(10)
    return (w1, b1, w2, b2)

def loss_fn(params, x, y):
    # 前向传播 + 交叉熵损失——与我们 NumPy 版本相同的数学
    w1, b1, w2, b2 = params
    h = jnp.maximum(0, x @ w1 + b1)       # 隐藏层 + ReLU
    logits = h @ w2 + b2                    # 输出层
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    return -jnp.mean(log_probs[jnp.arange(y.shape[0]), y])

def sgd_step(params, x, y, lr):
    # 一个训练步：计算梯度，更新权重
    grads = grad(loss_fn)(params, x, y)     # JAX 自动对 loss_fn 求导
    return tuple(p - lr * g for p, g in zip(params, grads))

# 5 个学习率、5 套权重，同时训练
lr_array = jnp.array([0.001, 0.01, 0.1, 1.0, 10.0])
keys = random.split(random.PRNGKey(42), len(lr_array))
batched_params = vmap(init_params)(keys)
batched_step = jit(vmap(sgd_step, in_axes=(0, None, None, 0)))
batched_params = batched_step(batched_params, x_batch, y_batch, lr_array)

另一种方法是让每个配置在独立进程中训练。下面的示例假设 Keras 和训练数据已准备好。请将它作为带有主程序入口保护的脚本运行，并根据每个进程的内存和 CPU 线程占用来选择进程数：

import multiprocessing as mp

def train_one_config(args):
    # 用一个学习率训练一个模型——在自己的进程中运行
    lr, X_train, y_train = args
    model = keras.Sequential([
        keras.layers.Dense(128, activation="relu", input_shape=(784,)),
        keras.layers.Dense(10, activation="softmax"),
    ])
    model.compile(optimizer=keras.optimizers.SGD(learning_rate=lr),
                  loss="sparse_categorical_crossentropy")
    history = model.fit(X_train, y_train, epochs=10, batch_size=32,
                        validation_split=0.2, verbose=0)
    return lr, history.history

if __name__ == "__main__":
    with mp.get_context("spawn").Pool(5) as pool:
        results = pool.map(train_one_config,
                           [(lr, X_train, y_train) for lr in [0.001, 0.01, 0.1, 1.0, 10.0]])

几点要带走的结论

1. 小模型可能在 CPU 上运行更快。最初的 Keras 测试中，CPU 每轮耗时 4.4 秒，GPU 为 6.5 秒；另一组批大小测试得到了不同的耗时。
2. 峰值 TFLOPS 不能用来衡量利用率。仅算术部分的耗时下限，无法告诉我们内核的实际耗时或活跃核心的数量。
3. 测量完整的工作负载。同步、验证、丢弃批次和编译都会影响计时结果的含义。
4. 比较相同条件下的训练。增大批次和编译循环可以减少耗时，但在宣称加速之前，还应比较验证效果。
5. 对手头的设备进行实测。仅凭参数数量，无法判断 CPU 和 GPU 哪个更快。

Word2vec也是一个说明数据访问方式为何重要的例子。采用负采样时，每次更新只访问嵌入矩阵中选定的行，而不是执行大型稠密矩阵乘法。CPU 和 GPU 实现的相对性能取决于这些查找和更新如何组织成批次，并不存在普遍适用的 CPU 优势。