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

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

猜你喜欢

转载自blog.csdn.net/u013289254/article/details/105278293