最近同事碰到了一个不同框架模型互相转换的问题：pytorch模型或caffe模型要转到tensorflow和TFLite上进行移动端的部署。模型从pytorch或caffe转tensorflow通过ONNX作为中转还相对顺利，但是从tensorflow到TFLite的时候却遇到了输入不匹配的问题。原因在于pytorch，caffe输入数据通常以“NCHW”的形式进行排列。虽然转到tensorflow的时候依旧也能按照“NCHW”的方式转成功，但是Tensorflow转TFLite竟然失败了。原来TFLite仅支持“NHWC”这种tensorflow系的数据排布方式。最后我同事虽然解决了，但是我也有了疑问：为啥各个框架的数据内存排布会不一样？

首先，在深度学习框架中，通常的卷积神经网络特征图用四维数组保存，四个维度分别为“N”batch批量大小，“C”channels特征图通道数，“H”特征图的高，和“W”特征图的宽。虽然我们人可以的将数据按照任意维度进行划分，然而对于计算机而言，数据的存储只能是线性的。假如一个特征图维度为[1，3，2，2]，实际上存储的是一块1x3x2x2=12个顺序的数据块。

在pytorch和caffe的排布方式中，一张2x2的RGB格式的图像，使用“NCHW”实际存储的是“RRRRGGGGBBBB”，同一个通道的所有像素值顺序存储在一起。而Tensorflow通常的排列使用“NHWC”为“RGBRGBRGBRGB”, 多通道的同一位置的像素值顺序存储在一起。

尽管存储的数据实际上是一样的，但是不同的顺序会导致数据的访问特性不一致，因此即使进行同样的运算，相应的计算性能也会不一样。对于"NCHW" 而言，其同一个通道的像素值连续排布，更适合那些需要对每个通道单独做运算的操作，比如"MaxPooling"。对于"NHWC"而言，其不同通道中的同一位置元素顺序存储，因此更适合那些需要对不同通道的同意像素做某种运算的操作，比如“Conv1x1”。如下图：

从上图可以看出，"NCHW"的计算时需要的存储更多，适合GPU运算，正好利用了GPU内存带宽较大并且并行性强的特点，其访存与计算的控制逻辑相对简单。而"NHWC"更适合多核CPU运算，CPU的内存带宽相对较小，每个像素计算的时延较低，临时空间也很小，有时计算机采取异步的方式边读边算来减小访存时间，因此计算控制灵活且复杂。因此，深度学习以cuDNN为底层的pytoch和caffe框架默认使用了 "NCHW" 格式，而Tensorflow采用了"NHWC"，据说是由于早期主要使用CPU加速，这也解释了为何面向移动端部署的TFLite只采用了"NHWC" 格式。

除了这两种排布方式，还有别的排布方式吗？ 答案是当然有。

比如说面向移动端的阿里出品的MNN框架，是以 "NC4HW4" 五维数据格式作为排布的，具体详情的可以参考：

Tensor中数据摆放顺序NC4HW4是什么意思，只知道NCHW格式，能解释以下NC4HW4格式吗？​www.zhihu.com

又比如华为海思的昇腾Ascend310芯片系列，其深度学习框架为了提高数据的访存效率，所有的张量数据统一也采用了 "NC1HWC0"的五维格式。由于其AI core的芯片设计，其cube矩阵计算单元的大小要求有要求，只能支持16x16的矩阵运算。其中对于FP16类型的数据要对齐16，对于INT8类型的数据要对齐32（其实与FP16等价）。其中五维度的"C0" 与这个数据对齐强相关，"C0, C1"其实就是从原本"NCHW"的"C"维度进行了拆分，"C1 = C / C0", 如果不能整除，则最后需要补零来对齐。其大致的拆解过程为：

（1）将"NCHW"数据在C维度进行分割，分出C1份NHWC0;

(2) 将C1份"NHWC0"在内存中连续排列，变成了"NC1HWC0"。

从阿里的MNN与华为的昇腾系列芯片的数据排布来看，不难发现他们都在"C"这个Channel维度做文章。看起十分复杂，但其实这非常合理和高效。由于典型的卷积神经网络随着层数的增加，其特征图在下采样后的长和宽逐渐减小，但是channel数随着卷积的filter的个数不断增大是越来越大的，经常会出现channel数为128，256等很深的特征图。这些很深的特征图与filter数很多的卷积层进行运算是很heavy的。为了充分利用有限的矩阵计算单元，进行了Channel维度的拆分是很有必要的。

那么有没有针对H和W维度进行优化的呢？答案也是当然有的。很多运算会关注数据时行优先（Row-Major）还是列优先（Column-Major）。矩阵运算的时候有左右矩阵之分，在这两个矩阵一个需要按行优先，一个需要按列优先访问。但是实际上输入的NCHW数据，HW会根据实际应用改变，比如1080P的输入数据，W比H大；Person-reid任务中图像H比W大。输入大小往往决定了特征图的输出大小。因此在输入矩阵前，往往会通过软件或者硬件对数据进行重排列，以达到最优的计算效率，正如上面所说的华为NCHW->NC1HWC0。

当然上述的排列方式没有孰优孰劣之分，不同的排布方式也有相应的底层优化算法，对上层用户来说这些差异是看不见的。深度学习引擎特别会在推理的时候，根据实际硬件结构和计算资源，对数据排布进行转换，有些中间过程的数据的排布和访存甚至会有多种优化方式。这也是很多kernel优化所要做的。