3、内存布局优化：NCHW与NHWC格式选择、内存对齐策略

内存布局这事儿，说白了就是数据在内存里怎么排排坐。很多刚入行的同学觉得这有啥好纠结的？不就是个维度顺序嘛。嗯，等你真正部署模型时就会发现，选错格式，性能直接腰斩。我当年第一次部署MobileNet时就吃过这个亏，后面细说。

3.1 NCHW vs NHWC：到底选哪个？

先看这两个格式长啥样。假设你有一张图片，batch=2，通道=3，高=4，宽=5：

• NCHW：先排通道，再排空间。内存里是 [batch0_ch0全部像素, batch0_ch1全部像素, ...]
• NHWC：先排空间，再排通道。内存里是 [batch0_pixel0的RGB, batch0_pixel1的RGB, ...]

你可能会问：这有啥区别？区别大了去了。我直接给结论：

特性	NCHW	NHWC
卷积计算	连续读取同一通道，cache友好	需要跨通道读取，cache命中率低
批量归一化	天然按通道计算，效率高	需要转置或特殊处理
逐像素操作	跨通道不连续，性能差	连续读取，性能好
硬件偏好	GPU（cuDNN默认）	CPU（ARM、x86优化库）

我记得有一次在Jetson Nano上部署一个分类模型，默认用NCHW格式跑，推理时间35ms。换成NHWC后，直接降到22ms。为什么？因为Jetson的GPU是Tegra架构，它对NHWC有硬件加速指令。你想想看，选错格式白白浪费30%的性能，多冤。

3.2 内存对齐：别小看这16字节

内存对齐是什么？简单说就是数据地址要能被某个数整除。最常见的是16字节对齐，也有32字节、64字节对齐。

为什么要对齐？因为CPU和GPU从内存读数据时，一次读一整块。如果数据没对齐，就得读两次再拼起来。我见过一个极端案例：某个算子因为没对齐，性能差了5倍。

来看个实际例子。假设你要分配一个float32的矩阵，宽=3，高=4：

// 不对齐的分配
float* data = (float*)malloc(3 * 4 * sizeof(float));
// 每行3个float = 12字节，不是16的倍数

// 对齐的分配
int pitch = ((3 * sizeof(float) + 15) / 16) * 16;  // 16字节对齐
float* data = (float*)aligned_alloc(16, pitch * 4);
// 每行16字节，虽然浪费了4字节，但访问效率高

你可能会觉得浪费空间。没错，确实浪费。但你要算笔账：浪费4字节换来的是每次内存访问都对齐，省下的时间远大于这点空间开销。我在项目中做过测试，对齐后卷积算子的速度能提升10%-30%。

3.3 实战中的格式转换策略

模型部署时，经常需要在不同格式间转换。我建议的策略是：

1. 训练阶段：用框架默认格式。PyTorch默认NCHW，TensorFlow默认NHWC。别折腾。
2. 导出阶段：根据目标硬件选择。比如导出到ONNX时，我习惯转成NCHW，因为ONNX Runtime对NCHW支持最好。
3. 推理阶段：如果硬件支持，尽量保持格式不变。频繁转换格式本身就有开销。

3.4 内存对齐的避坑指南

说几个我实际遇到过的坑：

• 跨平台对齐差异：x86上malloc默认16字节对齐，但ARM上不一定。我在树莓派上就遇到过这个问题，malloc返回的地址只8字节对齐，跑NEON指令直接崩。
• 结构体对齐：如果你自定义数据结构存张量，记得用__attribute__((aligned(16)))或alignas(16)。
• 共享内存对齐：多进程共享内存时，对齐要求更严格。我建议至少64字节对齐，避免不同进程的cache line冲突。

嗯，这里要注意：对齐不是万能的。如果你的张量很小（比如1x1卷积），对齐带来的收益微乎其微。我一般只在特征图尺寸大于32x32时才做对齐优化。

3.5 总结一下我的经验

内存布局优化，说白了就是两件事：选对格式、对齐地址。我个人的工作流是这样的：

• 拿到一个新模型，先看目标硬件文档，确定它偏好哪种格式
• 写一个简单的benchmark，对比NCHW和NHWC的性能差异
• 所有张量分配都做16字节对齐，除非内存极度受限
• 在关键路径上检查内存对齐情况，用(uintptr_t)ptr % 16 == 0来验证

你想想看，这些优化都不需要改模型结构，纯属工程技巧。但就是这些细节，决定了你的模型能不能跑到实时。我见过太多人花大量时间调模型结构，结果部署时因为内存布局不对，性能还不如没优化的版本。所以，先把这些基本功打好，再去搞那些花里胡哨的优化。