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

N: batch;

C: channel

H: height

W: width

Caffe 的Blob通道顺序是：NCHW;

Tensorflow的tensor通道顺序：默认是NHWC， 也支持NCHW，使用cuDNN会更快;

Pytorch中tensor的通道顺序：NCHW

TensorRT中的tensor 通道顺序： NCHW

numpy image: H x W x C
torch image: C X H X W

---

 

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

 

参考链接：https://blog.csdn.net/edward_zcl/article/details/102661617

原文链接：https://blog.csdn.net/u013289254/article/details/105278293

---

以pytorch为例，注意：

torch.nn只支持小批量处理(mini-batches）。整个torch.nn包只支持小批量样本的输入，不支持单个样本的输入。

比如，nn.Conv2d 接受一个4维的张量，即nSamples x nChannels x Height x Width

如果是一个单独的样本，只需要使用input.unsqueeze(0)来添加一个“假的”批大小维度。