内存布局优化：NHWC vs NCHW

说到内存布局，我刚开始做算子优化时，其实没太当回事。总觉得算法写对了就行，数据怎么排不都一样吗？直到有一次，我优化一个卷积算子，折腾了两周性能纹丝不动。后来一位老前辈看了一眼，说：「你把数据排布换一下试试。」结果换了之后，吞吐直接翻倍。嗯，从那以后，我再也不敢小看内存布局了。

两种主流数据排布格式

在深度学习框架里，最常见的就是两种数据排布：NCHW 和 NHWC。说白了，它们只是四个维度的排列顺序不同。

N：Batch Size，一次处理多少张图
C：Channel，通道数（比如RGB就是3通道）
H：Height，图片高度
W：Width，图片宽度

NCHW 把通道放在前面，NHWC 把通道放在最后。就这么一个简单的顺序差异，对性能的影响却天差地别。

核心区别：NCHW 是「通道连续」存储，NHWC 是「像素连续」存储。你想想看，CPU/GPU 访问内存时，最喜欢的就是连续地址。一旦跳来跳去，缓存命中率就直线下降。

数据排布对缓存命中率的影响

我习惯用一个简单的例子来说明这个问题。假设我们有一张 3×3 的 RGB 图片，每个像素点有 R、G、B 三个值。

NCHW 排布：

R R R R R R R R R  G G G G G G G G G  B B B B B B B B B
↑ 连续9个R值       ↑ 连续9个G值       ↑ 连续9个B值

NHWC 排布：

R G B  R G B  R G B  R G B  R G B  R G B  R G B  R G B  R G B
↑ 第一个像素       ↑ 第二个像素       ↑ 第三个像素...

为什么会这样？因为不同的算子对数据访问模式完全不同：

算子类型	更友好的排布	原因
卷积（Conv2D）	NCHW（GPU）/ NHWC（CPU）	GPU上NCHW利于向量化，CPU上NHWC利于缓存
批归一化（BatchNorm）	NCHW	需要对每个通道做统计，通道连续访问更快
逐元素操作（ReLU、Add）	两者均可	操作独立，但NHWC在CPU上缓存更友好
全连接（FC）	NCHW	需要将特征图展平，NCHW展平时更连续

我在项目中遇到过这样一个场景：一个 MobileNet 模型在 ARM CPU 上推理，默认用 NCHW 排布，帧率只有 15fps。我手动改成 NHWC 后，帧率直接飙到 28fps。原因很简单——MobileNet 大量使用深度可分离卷积，每个像素点的计算独立，NHWC 让数据访问完全连续，缓存命中率从 60% 提升到了 95% 以上。

我的经验：如果你在 CPU 上做推理，优先试试 NHWC。GPU 上则要看框架——TensorFlow 默认 NHWC，PyTorch 默认 NCHW。别盲目跟风，最好用 profiler 跑一下看看缓存 miss 率。

实战：手动调整Tensor布局

光说不练假把式。我们直接上手，手动实现一个 NCHW 到 NHWC 的转换。这个操作在深度学习里叫 「permute」 或者 「transpose」。

假设我们有一个形状为 [N, C, H, W] 的张量，要转成 [N, H, W, C]。说白了，就是把通道维度从第1位挪到最后1位。

// C++ 示例：NCHW -> NHWC 手动实现
void nchw_to_nhwc(const float* src, float* dst, 
                  int N, int C, int H, int W) {
    for (int n = 0; n < N; n++) {
        for (int h = 0; h < H; h++) {
            for (int w = 0; w < W; w++) {
                for (int c = 0; c < C; c++) {
                    // NCHW 索引: n*C*H*W + c*H*W + h*W + w
                    int src_idx = n*C*H*W + c*H*W + h*W + w;
                    // NHWC 索引: n*H*W*C + h*W*C + w*C + c
                    int dst_idx = n*H*W*C + h*W*C + w*C + c;
                    dst[dst_idx] = src[src_idx];
                }
            }
        }
    }
}

这段代码看起来简单，但性能其实很差。为什么？因为内层循环在 c 维度上，每次访问 src 时，c*H*W 这个跨度太大了，缓存根本扛不住。

我曾经踩过的坑：直接用上面这种四重循环做 transpose，处理一张 1080p 的图片（1920×1080×3）就要花 30 多毫秒。后来我改成用分块（tiling）和 SIMD 指令优化，才降到 2 毫秒以内。记住：内存搬运操作，一定要考虑缓存行大小（通常 64 字节）。

更好的做法是使用 分块转置：

// 优化版：分块转置，利用缓存局部性
void nchw_to_nhwc_blocked(const float* src, float* dst,
                           int N, int C, int H, int W) {
    const int BLOCK_H = 8;  // 分块大小，匹配缓存行
    const int BLOCK_W = 8;
    
    for (int n = 0; n < N; n++) {
        for (int h = 0; h < H; h += BLOCK_H) {
            for (int w = 0; w < W; w += BLOCK_W) {
                int h_end = min(h + BLOCK_H, H);
                int w_end = min(w + BLOCK_W, W);
                
                // 处理一个小块内的所有通道
                for (int bh = h; bh < h_end; bh++) {
                    for (int bw = w; bw < w_end; bw++) {
                        for (int c = 0; c < C; c++) {
                            int src_idx = n*C*H*W + c*H*W + bh*W + bw;
                            int dst_idx = n*H*W*C + bh*W*C + bw*C + c;
                            dst[dst_idx] = src[src_idx];
                        }
                    }
                }
            }
        }
    }
}

分块之后，内层循环访问的 src 数据都在一个连续的小区域内，缓存命中率大幅提升。我实测过，对于 224×224 的图片，分块版本比朴素版本快了 5-8 倍。

用 SVG 画一张内存布局对比图

下面这张图展示了 NCHW 和 NHWC 在内存中的实际存储方式，以及缓存访问的差异：

从这张图可以看得很清楚：NCHW 把同一个通道的所有像素放在一起，访问不同通道的同一个像素时，需要跳过一大段内存。而 NHWC 把同一个像素的所有通道放在一起，访问时完全连续。

我的建议：在实际项目中，不要手动写 transpose。直接用框架提供的 API——PyTorch 用 permute()，TensorFlow 用 tf.transpose()。这些底层实现已经用 CUDA/NEON 优化过了。但理解原理很重要，因为当你需要手写算子时，这个知识能帮你避免性能灾难。

最后说一句：内存布局没有绝对的好坏。我见过有人死磕 NHWC，结果在 GPU 上反而更慢。最好的做法是：先 profiler，后优化。用 perf 或者 nvprof 看看缓存 miss 率，再决定要不要改布局。这才是工程师该有的思路。

想要更多资料，联系可可

📚 课程目录

内存布局优化：NHWC vs NCHW

两种主流数据排布格式

数据排布对缓存命中率的影响

实战：手动调整Tensor布局

用 SVG 画一张内存布局对比图