流行的深度学习框架中有不同的数据格式,典型的有NCHW和NHWC格式。
下面从逻辑表达和物理存储角度用图的方式来理解这两种数据格式,最后以RGB图像为例来加深NHWC和NCHW数据存储格式的理解。
Intel(R) MKL-DNN: Understanding Memory Formats
深度学习框架中,数据一般是4D,用NCHW或NHWC表达,其中:
- N - Batch
- C - Channel
- H - Height
- W - Width
假定N = 2,C = 16,H = 5,W = 4,那么这个4D数据,看起来是这样的
但是无论逻辑表达上是几维的数据,在计算机中存储时都是按照1D来存储的。NCHW和NHWC格式数据的存储形式如下图所示:
NCHW是先取W方向数据;然后H方向;再C方向;最后N方向。
所以,序列化出1D数据:
000 (W方向) 001 002 003,(H方向) 004 005 ... 019,(C方向) 020 ... 318 319,(N方向) 320 321 ...
NHWC是先取C方向数据;然后W方向;再H方向;最后N方向。
所以,序列化出1D数据:
000 (C方向) 020 ... 300,(W方向) 001 021 ... 303,(H方向) 004 ... 319,(N方向) 320 340 ...
表达RGB彩色图像时,一个像素的RGB值用3个数值表示,对应Channel为3。易于理解这里假定N=1,那么NCHW和NHWC数据格式可以很直接的这样表达:
NCHW是先在一个Channel面上把W方向|H方向上元素存储起来 // R
然后再在另一个Channel切面上把W方向|H方向上元素存储起来 // G
最后一个Channel切面上把W方向|H方向上元素存储起来 // B
这样看起来,就是先把R通道的每个像素都存储;然后存储G通道;再然后B通道。
NHWC是先把3个Channel上元素存储起来 // 也就是一个像素的RGB
然后再在W方向|H方向上元素存储起来
这样看起来,就是顺序地取像素的RGB数值存储起来
目前的主流ML框架对NCHW和NHWC数据格式做了支持,有些框架可以支持两种且用户未做设置时有一个缺省值:
- TensorFlow:缺省NHWC,GPU也支持NCHW
- Caffe:NCHW
- PyTorch:NCHW