Caffe、Tensorflow和Pytorch通道维度顺序小结

 目前世面上深度学习框架比较多，常用的有三大类：Caffe、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

参考链接：TensorFlow API 中 NCHW 与 NHWC 的区别_edward_zcl的博客-CSDN博客