一、参考资料
图解NCHW与NHWC数据格式
TensorFlow API 中 NCHW 与 NHWC 的区别
深度学习中的NCHW格式和NHWC格式
二、重要说明
- 在 GPU 上训练时使用 NCHW 格式,在 CPU 上做预测时使用 NHWC 格式。NCHW 则是 Nvidia cuDNN 默认格式,使用 GPU 加速时用 NCHW 格式速度会更快(也有个别情况例外)。
- 最佳实践:设计网络时充分考虑两种格式,最好能灵活切换,在 GPU 上训练时使用 NCHW 格式,在 CPU 上做预测时使用 NHWC 格式。
三、相关介绍
3.1 NCHW的含义
- N - Batch
- C - Channel
- H - Height
- W - Width
3.2 逻辑表达
假定N = 2,C = 16,H = 5,W = 4,那么这个4D数据如下图所示:
人类比较直接的理解方式是3D,上图中从三个方向上理解,C方向/H方向/W方向。然后是N方向上,就是4D。上图中红色标准的数值是这个数据里每个元素的数值。
3.3 物理存储
无论逻辑表达上是几维的数据,在计算机中存储时都是按照1D来存储的。
3.3.1 NCHW
NCHW是先取W方向数据;然后H方向;再C方向;最后N方向。
所以,序列化出1D数据:
000 (W方向) 001 002 003,(H方向) 004 005 … 019,(C方向) 020 … 318 319,(N方向) 320 321 …
3.3.2 NHWC
NHWC是先取C方向数据;然后W方向;再H方向;最后N方向。
所以,序列化出1D数据:
000 (C方向) 020 … 300,(W方向) 001 021 … 303,(H方向) 004 … 319,(N方向) 320 340 …
3.4 RGB图像数据举例
表达RGB彩色图像时,一个像素的RGB值用3个数值表示,对应Channel为3。易于理解这里假定N=1,那么NCHW和NHWC数据格式可以很直接的这样表达:
3.4.1 NCHW
- 先在一个Channel面上把W方向|H方向上元素存储起来 // R
- 然后再在另一个Channel切面上把W方向|H方向上元素存储起来 // G
- 最后一个Channel切面上把W方向|H方向上元素存储起来 // B
这样看起来,就是先把R通道的每个像素都存储;然后存储G通道;再然后B通道。
3.4.2 NHWC
- 先把3个Channel上元素存储起来 // 也就是一个像素的RGB
- 然后再在W方向|H方向上元素存储起来
这样看起来,就是顺序地取像素的RGB数值存储起来。
四、不同框架支持
目前的主流ML框架对NCHW和NHWC数据格式做了支持,有些框架可以支持两种且用户未做设置时有一个缺省值:
-
TensorFlow:缺省NHWC,GPU也支持NCHW;
TensorFlow 为什么选择 NHWC 格式作为默认格式?因为早期开发都是基于 CPU,使用 NHWC 比 NCHW 稍快一些(不难理解,NHWC 局部性更好,cache 利用率高)。
-
Caffe:NCHW;
-
PyTorch:NCHW;
五、NCHW与NHWC转换
# NCHW [batch,in_channels,in_height,in_weight]
# NHWC [batch,in_height,in_weight,in_channels]
# CHWN [in_channels,in_height,in_weight,batch]
# 转换 NCHW---NHWC
import tensorflow as tf
x = tf.reshape(tf.range(24),[1,2,3,4])
out = tf.transpose(x,[0,2,3,1])
print (x.shape)
print (out.shape)
#转换NHWC--NCHW
import tensorflow as tf
x = tf.reshape(tf.range(24), [1, 3, 4, 2])
out = tf.transpose(x, [0, 3, 1, 2])
print (x.shape)
print (out.shape)