Tensorflow和Pytorch通道维度总结

Caffe 的通道顺序是NCHW;

Tensorflow的通道顺序默认是NHWC（但可以设置成NCHW），NHWC 的访存局部性更好（每三个输入像素即可得到一个输出像素），NCHW 则必须等所有通道输入准备好才能得到最终输出结果，需要占用较大的临时空间。

TensorFlow 为什么选择 NHWC 格式作为默认格式？因为早期开发都是基于 CPU，使用 NHWC 比 NCHW 稍快一些（不难理解，NHWC 局部性更好，cache 利用率高）。

NCHW 则是 Nvidia cuDNN 默认格式，使用 GPU 加速时用 NCHW 格式速度会更快（也有个别情况例外）。

最佳实践：设计网络时充分考虑两种格式，最好能灵活切换，在 GPU 上训练时使用 NCHW 格式，在 CPU 上做预测时使用 NHWC 格式。

Pytorch的通道顺序是NCHW

NCHW 中，C 排列在外层，每个通道内像素紧挨在一起，即 ‘RRRRRRGGGGGGBBBBBB’ 这种形式。
NHWC 格式，C 排列在最内层，多个通道对应空间位置的像素紧挨在一起，即 ‘RGBRGBRGBRGBRGBRGB’ 这种形式。
如果我们需要对图像做彩色转灰度计算，NCHW 计算过程如下：
在这里插入图片描述

即 R 通道所有像素值乘以 0.299，G 通道所有像素值乘以 0.587，B 通道所有像素值乘以 0.114，最后将三个通道结果相加得到灰度值。
相应地，NHWC 数据格式的彩色转灰度计算过程如下：
在这里插入图片描述

输入数据分成多个(R, G, B) 像素组，每个像素组中 R 通道像素值乘以 0.299，G 通道像素值乘以 0.587，B 通道像素值乘以 0.114 后相加得到一个灰度输出像素。将多组结果拼接起来得到所有灰度输出像素。
以上使用两种数据格式进行 RGB -> 灰度计算的复杂度是相同的，区别在于访存特性。通过两张图对比可以发现，NHWC 的访存局部性更好（每三个输入像素即可得到一个输出像素），NCHW 则必须等所有通道输入准备好才能得到最终输出结果，需要占用较大的临时空间。