3、内存布局优化：理解NCHW vs NHWC，数据排布对访存效率的影响

内存布局这事儿，说白了就是数据在内存里怎么排队。

很多同学刚开始写算子，觉得只要算对了就行。嗯，我当年也这么想。直到有一次，我写了个卷积算子，理论上FLOPs算得明明白白，结果跑起来比PyTorch慢了20倍。查了两天，最后发现——问题出在数据怎么放，而不是怎么算。

今天我们就聊聊NCHW和NHWC这两种排布方式。它们不复杂，但搞错了，性能差距能到数量级。

3.1 两种排布方式：一张图说清楚

先看个直观的例子。假设我们有一张图片：

• 批次大小 Batch = 2
• 通道数 Channel = 3（RGB）
• 高度 Height = 4
• 宽度 Width = 4

数据总量是 2×3×4×4 = 96 个元素。但怎么排到一维内存里？

NCHW：先按通道排，再按行排。也就是：

Batch 0:
  Channel 0: [所有行所有列]
  Channel 1: [所有行所有列]
  Channel 2: [所有行所有列]
Batch 1:
  Channel 0: [所有行所有列]
  ...

NHWC：先按行排，再按通道排。也就是：

Batch 0:
  Row 0: [C0, C1, C2, C0, C1, C2, ...]
  Row 1: [C0, C1, C2, ...]
  ...
Batch 1:
  ...

我画了张图，帮你直观理解这两种排布的内存访问模式：

3.2 为什么排布方式影响性能？

原因就三个字：缓存行。

CPU/GPU从内存读数据，不是一次读一个数，而是一次读一整块（比如64字节）。如果你要的数据刚好在同一个缓存行里，那就快；如果分散在不同缓存行，那就慢。

举个例子。假设你在做卷积，卷积核是3×3：

• NCHW 下：要取某个通道的3×3区域，这9个数在内存里是连续的（同一通道的相邻行）。一次缓存行读取，可能就把9个数全拿回来了。快。
• NHWC 下：要取某个通道的3×3区域，但数据是按像素交错存放的。你读一个R，中间夹着G和B，再读下一个R，中间又夹着G和B。缓存行里有一半数据是你不需要的。浪费。

反过来，如果做逐像素操作（比如归一化、激活函数），NHWC 反而更优。因为你要同时处理同一个像素的R、G、B，它们正好挨在一起。

3.3 不同框架的默认选择

框架	默认布局	备注
PyTorch	NCHW	图像输入默认是 (B, C, H, W)
TensorFlow	NHWC	历史原因，早期CPU上NHWC更快
ONNX	NCHW	标准中间表示
TensorRT	NCHW	但内部会做布局优化

你想想看，如果你用PyTorch训练，导出ONNX，再用TensorRT推理，中间其实经历了一次隐式的布局转换。这个转换本身也有开销。

3.4 实战：如何手动转换布局

有时候框架不帮你做转换，你得自己来。比如你写了一个自定义算子，输入是NHWC，但你的计算逻辑需要NCHW。

在Python里，用NumPy做转换很简单：

import numpy as np

# 假设输入是 NHWC 格式
x_nhwc = np.random.randn(2, 4, 4, 3)  # (B, H, W, C)

# 转换为 NCHW
x_nchw = np.transpose(x_nhwc, (0, 3, 1, 2))  # (B, C, H, W)

# 再转回 NHWC
x_nhwc_again = np.transpose(x_nchw, (0, 2, 3, 1))

在C++里，如果你用TensorRT的插件，可以这样：

// NHWC -> NCHW
// 假设输入是 float* src, 输出是 float* dst
// dims: [N, H, W, C]
for (int n = 0; n < N; ++n) {
    for (int c = 0; c < C; ++c) {
        for (int h = 0; h < H; ++h) {
            for (int w = 0; w < W; ++w) {
                // NHWC: index = n*H*W*C + h*W*C + w*C + c
                // NCHW: index = n*C*H*W + c*H*W + h*W + w
                dst[n*C*H*W + c*H*W + h*W + w] = src[n*H*W*C + h*W*C + w*C + c];
            }
        }
    }
}

3.5 硬件偏好：CPU vs GPU

不同的硬件，对布局的偏好也不同：

• CPU：偏向NHWC。因为CPU的SIMD指令（如AVX）擅长处理连续内存，而NHWC在逐像素操作时数据更紧凑。
• GPU（NVIDIA）：偏向NCHW。因为cuDNN对NCHW做了深度优化，卷积操作时能更好地利用共享内存。
• ARM/Mobile：偏向NHWC。移动端的NPU和DSP通常对NHWC更友好。

我记得有一次在Jetson Nano上做部署，默认用NCHW，结果帧率只有15fps。改成NHWC后，直接到了30fps。嗯，硬件差异就是这么明显。

3.6 避坑指南：我踩过的三个坑

坑一：忽略布局转换的代价

我曾经写过一个算子，输入是NHWC，内部计算需要NCHW。我每调用一次算子就做一次转换。结果性能比PyTorch慢了10倍。后来改成在算子外部统一转换，内部直接用NHWC计算（虽然计算复杂了点），性能反而上来了。

坑二：以为所有框架都一样

有一次我把PyTorch的模型直接搬到TensorFlow Serving上，没改布局。结果推理结果全是错的。查了半天，发现PyTorch默认NCHW，TensorFlow默认NHWC，数据解析就错了。

坑三：忽略通道数的影响

当通道数很少时（比如3通道RGB），NHWC的劣势不明显。但当通道数很多时（比如ResNet的512通道），NCHW的优势就非常明显了。因为连续的大块内存访问，对缓存更友好。