以下内容翻译自: Understanding Memory Formats
Introduction
大多数计算都是关于数据的:分析数据、调整数据、读取和存储数据、生成数据等。DNN 领域也不例外。图像、权重/过滤器、声音和文本需要在计算机内存中高效表示,从而以最方便的方式快速执行操作。
本文致力于数据格式的一种数据表示形式,它描述了多维数组(nD)如何存储在线性(1D)内存地址空间中,以及为什么这对 oneDNN 很重要。
注意: 在本文中,数据格式和布局可以互换使用。
使用的术语
- 通道与特征图相同
- 大写字母表示维度(例如
N
) - 小写字母表示索引(例如
n
,其中0<=n<n
) - 激活的符号:批次
N
,通道C
,深度D
,高度H
,宽度W
- 权重的记法:组别
G
,输出通道O
,输入通道I
,深度D
,高度H
,宽度W
Data Formats
让我们首先关注激活(图像)的数据格式。
激活由通道(也称为特征图)和空间域(1D、2D 或3D)组成。空间域与通道一起构成图像。 在训练阶段,图像通常按批次分组。即使只有一幅图像,我们仍然假设存在一个批大小等于1的批量。因此,激活的整体维度是4D(N、C、H 和 W)或5D(N、C、D、H 和 W)。
为了简单起见,本文将只使用2D 空间。
普通数据布局
从一个例子开始会更简单。
考虑批次等于2、16个通道和5 x 4空间域的4D 激活。下图给出了逻辑表示。
位置(n, c, h, w)
处的值由下式生成:
value(n, c, h, w) = n * CHW + c * HW + h * W + w
为了定义这个4D 张量中的数据是如何在内存中布局的,我们需要定义如何通过一个偏移量函数将其映射到一个1D 张量,该函数以逻辑索引(n, c, h, w)
作为输入,并返回一个到值所在位置的地址位移:
offset : (int, int, int, int) --> int
NCHW
让我们来描述一种非常流行的格式——NCHW,其张量值在内存中的排列顺序。[a:?]
标记是指下图所示的跳跃,它表示 NCHW 张量在内存中的一维表示。
[a:0]
:一行中的第一个,从左到右[a:1]
:然后自上而下逐行[a:2]
:然后从一个平面转到另一个平面(深度)[a:3]
:最后从一个批次(n=0)中的一个图像切换到另一个(n=1)
则偏移函数为:
offset_nchw(n, c, h, w) = n * CHW + c * HW + h * W + w
这里用nchw
表示w
是最内层的维度,这意味着内存中相邻的两个元素具有相同的n
、c
和h
索引,并且它们的w
索引相差1。当然,这仅适用于非边界元素。相反,n
是这里最外层的维度,这意味着如果需要在下一张图像上取相同的像素(c, h, w)
,则必须跳过整个图像大小C*H*W
。
这种数据格式称为 NCHW,在 BVLC Caffe 中默认使用。TensorFlow 也支持这种数据格式。
注意: 在本例中,
offset_nchw()
与value()
相同只是巧合。
对于 C API,可以使用 dnnl_nchw(dnnl_types.h 中定义的 dnnl_format_tag_t 枚举类型)创建具有 NCHW 数据布局的内存;C++ API 使用 dnnl.hpp 中定义的 dnnl::memory::format_tag::nchw 。
NHWC
另一种比流行的数据格式是 NHWC,它使用以下偏移函数:
offset_nhwc(n, c, h, w) = n * HWC + h * WC + w * C + c
在这种情况下,最内层的维度是通道([b:0]
),然后是宽度([b:1]
)、高度([b:2]
),最后是批次([b:3]
)。
对于单幅图像(N=1
),该格式与 BMP 文件格式的工作原理非常相似,其中图像逐像素保存,每个像素都包含所有需要的颜色信息(例如, 24位 BMP 的3个通道)。
NHWC数据格式是 TensorFlow 的默认格式。
此布局对应于 dnnl_nhwc 或 dnnl::memory::format_tag::nhwc。
CHWN
最后一个普通数据布局的例子是 Neon 使用的 CHWN。如果使用合适的批量大小,从向量化的角度来看,这种布局可能是非常有趣的,但另一方面,用户不可能总是拥有良好的批量大小(例如,当实时推理批处理通常为1)。
维度顺序为(从最内层到最外层):批次([c:0]
)、宽度([c:1]
)、高度([c:2]
)、通道([c:3]
)。
CHWN 格式的偏移函数定义为:
offset_chwn(n, c, h, w) = c * HWN + h * WN + w * N + n
该布局对应于 dnnl_chwn 或 dnnl::memory::format_tag::chwn。
Relevant Reading
TensorFlow Doc. Shapes and Layout
普通数据布局的推广
步长
在前面的示例中,数据以打包或以稠密的形式保存,这意味着像素彼此跟随。有时可能需要在内存中不保持数据连续。例如,有人可能需要在更大的张量中使用子张量。有时,人为地使数据不相交可能是有益的,例如GEMM采用非平凡的前导维数可以获得更好的性能(参见提示6)。
下图显示了以行主格式保存的大小为rows x columns
的二维矩阵的简化情况,其中行具有一些非平凡的(即,不等于列数)步长。
在这种情况下,一般的偏移函数如下所示:
offset(n, c, h, w) = n * stride_n
+ c * stride_c
+ h * stride_h
+ w * stride_w
需要注意的是,NCHW、NHWC 和 CHWN 格式只是带步长格式的特例。例如,对于 NCHW,我们有:
stride_n = CHW, stride_c = HW, stride_h = W, stride_w = 1
用户可以用步长初始化内存描述符:
dnnl_dims_t dims = {N, C, H, W};
dnnl_dims_t strides = {stride_n, stride_c, stride_h, stride_w};
dnnl_memory_desc_t md;
dnnl_memory_desc_init_by_strides(&md, 4, dims, dnnl_f32, strides);
oneDNN通过分块结构支持跨步。上述函数的伪代码为:
dnnl_memory_desc_t md; // memory descriptor object
// logical description, layout independent
int ndims = 4; // # dimensions
dnnl_dims_t dims = {N, C, H, W}; // dimensions themselves
dnnl_dims_t strides = {stride_n, stride_c, stride_h, stride_w};
dnnl_memory_desc_create_with_strides(&md, ndims, dims, dnnl_f32, strides);
特别地,每当用户以 dnnl_nchw 格式创建内存时,oneDNN 为用户计算步长并填充结构体。
Blocked Layout
平面布局提供了极大的灵活性,并且使用起来非常方便。这就是为什么大多数框架和应用程序使用 NCHW 或 NHWC 布局的原因。然而,根据对数据执行的操作,从性能角度来看,这些布局可能是次优的。
为了实现更好的向量化和缓存重用,oneDNN 引入分块布局,将一个或几个维度拆分成固定大小的块。AVX512+系统上最流行的 oneDNN 数据格式是 nChw16c,SSE4.1+ 系统上为 nChw8c。正如人们可能从名称中猜测的那样,仅对通道维度分块,并且块大小在前一种情况下为16,在后一种情况中为8。
准确来说,nChw8c 的偏移函数为:
offset_nChw8c(n, c, h, w) = n * CHW
+ (c / 8) * HW*8
+ h * W*8
+ w * 8
+ (c % 8)
注意,块中8个通道在内存中保持连续。逐像素覆盖空间域。然后,下一个切片覆盖随后的8个通道(即,从c=0..7
移动到c=8..15
)。 覆盖所有通道块后,将显示批处理中的下一个图像。
注意: 我们在格式中使用小写和大写字母来区分块(如8c)和剩余的联合维度(C=通道/8)。
格式选择背后的原因可以在 Distributed Deep Learning Using Synchronous Stochastic Gradient Descent 中找到。
oneDNN 也通过块结构来描述这种类型的内存。伪代码为:
dnnl_memory_desc_t md; // memory descriptor object
// logical description, layout independent
int ndims = 4; // # dimensions
dnnl_dims_t dims = {N, C, H, W}; // dimensions themselves
dnnl_memory_desc_create_with_tag(&md, ndims, dims, dnnl_f32, dnnl_nChw8c);
ptrdiff_t stride_n = C*H*W;
ptrdiff_t stride_C = H*W*8;
ptrdiff_t stride_h = W*8;
ptrdiff_t stride_w = 8;
dnnl_dims_t strides = {stride_n, stride_C, stride_h, stride_w }; // strides between blocks
int inner_nblks = 1; // number of blocked dimensions;
// 1, since only channels are blocked
dnnl_dims_t inner_idxs = {1}; // Only the 1st (c) dimension is blocked
// n -- 0st dim, w -- 3rd dim
dnnl_dims_t inner_blks = {8}; // This 1st dimensions is blocked by 8
dnnl_dims_t *q_strides = nullptr;
int *q_inner_nblks = nullptr;
dnnl_dims_t *q_inner_idxs = nullptr;
dnnl_dims_t *q_inner_blks = nullptr;
dnnl_memory_desc_query(md, dnnl_query_strides, &q_strides);
dnnl_memory_desc_query(md, dnnl_query_inner_nblks, &q_inner_nblks);
dnnl_memory_desc_query(md, dnnl_query_inner_idxs, &q_inner_idxs);
dnnl_memory_desc_query(md, dnnl_query_inner_blks, &q_inner_blks);
assert(memcmp(*q_strides, strides, DNNL_MAX_NDIMS) == 0);
assert(*q_inner_nblks == inner_nblks);
assert(memcmp(*q_inner_idxs, inner_idxs, DNNL_MAX_NDIMS) == 0);
assert(memcmp(*q_inner_blks, inner_blks, DNNL_MAX_NDIMS) == 0);
如果通道不是8(或16)的倍数呢?
分块数据布局给卷积带来了显著的性能提升,但当通道数不是块大小(例如,nChw8c 格式的17个通道)的倍数时该怎么办?
一种可能的处理方法是对尽可能多的通道使用分块布局,将它们四舍五入到一个块大小(此时16 = 17 / 8 * 8
)的倍数,并以某种方式处理尾部。然而,这将导致在许多 oneDNN 内核中引入非常特殊的尾部处理代码。
因此我们提出了另一种使用补零的解决方案。其思想是将通道舍入为块大小的倍数,并用零填充生成的尾部(在上面的示例中,24 = div_up(17, 8) * 8
)。然后,像卷积这样的原语可以使用四舍五入的通道数而不是原始通道数,并计算出正确的结果(添加零不改变结果)。
这使得可以在几乎不改变内核的情况下支持任意数量的通道。代价是在这些零点上进行一些额外的计算,但是要么这可以忽略不计,要么具有开销的性能仍然高于普通数据布局的性能。
下图描述了这个想法。注意,在d0
的计算过程中会产生一些额外的计算,但这并不会影响结果。
所给方法的一些缺陷:
- 保存数据所需的内存大小不能再通过公式
sizeof(data_type) * N * C * H * W
计算。实际大小应始终通过 C 中的 dnnl_memory_desc_get_size() 和 C++中的 dnnl::memory::desc::get_size() 查询。 - oneDNN 内存对象的实际补零发生在原语执行函数内部,以最小化其性能影响。目前的惯例是,一个原语执行可以假设其输入经过适当的零填充,并且应该保证其输出正确补零。如果用户在 oneDNN 分块内存对象上实现自定义内核,那么他们应该遵守这个约定。特别地,像这样在用户代码中实现并直接在 oneDNN 分块布局上运行的逐元素操作:
for (int e = 0; e < phys_size; ++e)
x[e] = eltwise_op(x[e])
若数据经过补零,且eltwise_op(0) != 0
,则不安全。
相关的 oneDNN 代码:
const int block_size = 8;
const int C = 17;
const int C_padded = div_up(17, block_size) * block_size;
const int ndims = 4;
memory::dims dims = {N, C, H, W};
memory::desc(dims, memory::data_type::f32, memory::format_tag::nChw8c);
memory::dim expect_stride_n = C_padded * H * W;
memory::dim expect_stride_C = H * W * block_size;
memory::dim expect_stride_h = W * block_size;
memory::dim expect_stride_w = block_size;
memory::dim expect_stride_8c = 1;
const bool expect_true = true
&& true // logical dims stay as is
&& md.get_dims()[0] == N
&& md.get_dims()[1] == C
&& md.get_dims()[2] == H
&& md.get_dims()[3] == W
&& true // padded dims are rounded accordingly
&& md.get_padded_dims()[0] == N
&& md.get_padded_dims()[1] == C_padded
&& md.get_padded_dims()[2] == H
&& md.get_padded_dims()[3] == W
&& true // strides between blocks correspond to the physical layout
&& md.get_strides()[0] == expect_stride_n
&& md.get_strides()[1] == expect_stride_C
&& md.get_strides()[2] == expect_stride_h
&& md.get_strides()[3] == expect_stride_w
&& true // inner-most blocking
&& md.get_inner_nblks() == 1 // only 1 dim is blocked (c)
&& md.get_inner_idxs()[0] == 1 // 1st (c) dim is blocked
&& md.get_inner_blks()[0] == 8; // the block size is 8
assert(expect_true);
参考资料:
- Memory Format Propagation
- 深度学习中的Tensor 数据格式(N,C,H,W)
- 深度学习工程化(1)多维图片在内存的高效存储和访问
- Distributed Deep Learning Using Synchronous Stochastic Gradient Descent
- Maximize TensorFlow* Performance on CPU: Considerations and Recommendations for Inference Workloads
- 最大限度提升 CPU 上的 TensorFlow* 性能:推理工作负载的注意事项和建议
- Understanding Memory Formats