Understanding Memory Formats

以下内容翻译自： Understanding Memory Formats

Introduction

大多数计算都是关于数据的：分析数据、调整数据、读取和存储数据、生成数据等。DNN 领域也不例外。图像、权重/过滤器、声音和文本需要在计算机内存中高效表示，从而以最方便的方式快速执行操作。

本文致力于数据格式的一种数据表示形式，它描述了多维数组（nD）如何存储在线性（1D）内存地址空间中，以及为什么这对 oneDNN 很重要。

注意： 在本文中，数据格式和布局可以互换使用。

使用的术语

• 通道与特征图相同
• 大写字母表示维度（例如N）
• 小写字母表示索引（例如n，其中0<=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