谈谈深度学习框架的数据排布【转载】

https://zhuanlan.zhihu.com/p/149464086?from_voters_page=true

最近同事碰到了一个不同框架模型互相转换的问题：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(上) 与NHWC(下)

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

NCHW MaxPooling (左)， NHWC Conv1x1 (右)

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