4、张量系统设计：张量数据结构定义、内存管理策略、数据布局（NHWC vs NCHW）
好，我们进入第四章。张量系统设计。
说实话，这是整个推理引擎的「地基」。你想想看，后面所有的算子、图优化、内存池，全都在这个数据结构上跑。地基没打好，后面全是坑。我自己最早做推理引擎时，就吃过这个亏——张量设计得太随意，结果后面加量化、加混合精度时，改得我头皮发麻。
4.1 张量数据结构定义
张量是什么？说白了，就是一个多维数组，外加一些元信息。但具体怎么定义，这里面有讲究。
我个人习惯把张量结构拆成三块：
- 形状信息：维度、步长、元素个数
- 数据信息：数据指针、数据类型、数据布局
- 设备信息：在哪个设备上（CPU/GPU）、设备ID
嗯，这里要注意：千万别把形状和步长混为一谈。我见过不少新手把形状当步长用，结果做切片时维度全乱了。
来看一个我常用的定义：
typedef struct {
int32_t ndim; // 维度数量
int32_t* dims; // 各维度大小，比如 [1, 3, 224, 224]
int32_t* strides; // 各维度步长，单位是元素个数
int64_t num_elements; // 总元素数
void* data; // 数据指针
DataType dtype; // 数据类型：float32, int8, uint8...
DataLayout layout; // 数据布局：NHWC 或 NCHW
DeviceType device; // 设备类型：CPU, CUDA, OpenCL...
int32_t device_id; // 设备编号
// 引用计数，用于内存管理
int32_t ref_count;
} Tensor;这里我特别想强调一下 strides。为什么要有它？因为推理引擎里经常要做「视图变换」——比如把一个 4D 张量 reshape 成 2D，或者做 transpose。如果没有 strides，每次都要重新分配内存，性能直接崩掉。
核心原则：张量的数据定义，要支持「零拷贝」的视图操作。这是高性能推理引擎的基石。
4.2 内存管理策略
内存管理，是推理引擎里最容易出 bug 的地方。我曾经在一个项目里，因为内存泄漏排查了整整三天，最后发现是张量析构时没释放设备内存。
我推荐三种策略，按场景选：
4.2.1 引用计数（Reference Counting）
这是最基础的做法。每个张量维护一个 ref_count，每被引用一次就加 1，释放时减 1。减到 0 就真正释放内存。
优点：简单，适合小规模推理。缺点：循环引用时会有泄漏风险。不过推理引擎里很少出现循环引用，所以够用了。
4.2.2 内存池（Memory Pool）
推理引擎里，张量频繁创建和销毁。每次都 malloc/free，性能太差。我建议用内存池：
- 预分配一大块内存
- 张量创建时从池里取
- 张量销毁时还回池里
- 池里按大小分桶，减少碎片
嗯，这里有个坑：GPU 内存池比 CPU 内存池难做得多。因为 GPU 内存分配本身就很慢，而且显存碎片更严重。我建议 GPU 内存池用「伙伴算法」或者「 slab 分配器」。
4.2.3 延迟分配（Lazy Allocation）
这个技巧我特别喜欢。很多张量在创建时，其实还不知道最终大小。比如动态 shape 的场景。这时候可以先分配一个「空壳」，等真正需要数据时再分配内存。
我的经验：在 ONNX 模型里，很多中间张量的 shape 是动态的。用延迟分配，可以避免大量无效的内存分配。我曾经在一个 BERT 模型上，用这个策略减少了 40% 的内存峰值。
4.3 数据布局：NHWC vs NCHW
好，到了最经典的话题。NHWC 还是 NCHW？
说白了，就是数据在内存里怎么排的。N 是 batch，C 是通道，H 是高度，W 是宽度。
- NCHW：先排所有通道的第一个像素，再排第二个像素……
- NHWC：先排第一个像素的所有通道，再排第二个像素的所有通道……
举个例子，一张 2x2 的 RGB 图片：
// NCHW 内存布局：[R00, G00, B00, R01, G01, B01, R10, G10, B10, R11, G11, B11]
// NHWC 内存布局：[R00, R01, R10, R11, G00, G01, G10, G11, B00, B01, B10, B11]你可能会问：哪个更好？答案是——看硬件。
| 布局 | 优点 | 缺点 | 典型场景 |
|---|---|---|---|
| NCHW | 通道内连续，适合卷积计算 | 跨通道访问不连续 | GPU（cuDNN 默认） |
| NHWC | 像素内连续，适合逐元素操作 | 卷积时需转置 | CPU（ARM、x86）、TensorFlow |
我个人在项目中遇到过这样一个情况：一个 MobileNet 模型，在 GPU 上用 NCHW 跑得飞快，但部署到手机端（ARM CPU）时，性能直接掉了 3 倍。后来改成 NHWC，性能才恢复正常。
避坑指南：我曾经在一个项目里，模型训练用 NCHW，推理引擎默认 NHWC，结果没做 layout 转换，跑出来的结果全是错的。嗯，这个坑我踩过，你别踩。
4.3.1 如何选择？
我的建议是：
- GPU 推理：用 NCHW。cuDNN 和 TensorRT 对 NCHW 优化最好。
- CPU 推理：用 NHWC。ARM 的 NEON 指令集对 NHWC 更友好。
- 混合场景：在引擎内部统一用一种布局，在边界处做 layout 转换。
嗯，这里要注意：layout 转换本身也有开销。我建议在模型加载时一次性转换，不要在推理过程中反复转。
4.3.2 扩展：其他布局
除了 NHWC 和 NCHW，还有一些变种：
- CHWN：Caffe 用的，现在已经很少见了
- NC4HW4：Intel 的 oneDNN 用的，把 C 维度拆成 4 的倍数
- NCHW16C：华为昇腾用的，16 通道对齐
这些变种说白了，都是为了适配特定硬件的 SIMD 指令宽度。你如果做通用推理引擎，建议先支持 NHWC 和 NCHW，其他的按需扩展。
总结一下：张量系统设计，核心就三件事——数据结构要灵活、内存管理要高效、数据布局要适配硬件。这三件事做好了，后面的算子实现会轻松很多。
好，这一章就到这里。下一章我们讲算子注册与调度，那是让张量「动起来」的关键。