Understanding Memory Formats

以下内容翻译自: 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 激活。下图给出了逻辑表示。
mem_fmt_img1

位置(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是最内层的维度,这意味着内存中相邻的两个元素具有相同的nch索引,并且它们的w索引相差1。当然,这仅适用于非边界元素。相反,n是这里最外层的维度,这意味着如果需要在下一张图像上取相同的像素(c, h, w),则必须跳过整个图像大小C*H*W

这种数据格式称为 NCHW,在 BVLC Caffe 中默认使用。TensorFlow 也支持这种数据格式。

注意: 在本例中,offset_nchw()value()相同只是巧合。

对于 C API,可以使用 dnnl_nchwdnnl_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_nhwcdnnl::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_chwndnnl::memory::format_tag::chwn
mem_fmt_img2

Relevant Reading

TensorFlow Doc. Shapes and Layout

普通数据布局的推广

步长

在前面的示例中,数据以打包或以稠密的形式保存,这意味着像素彼此跟随。有时可能需要在内存中不保持数据连续。例如,有人可能需要在更大的张量中使用子张量。有时,人为地使数据不相交可能是有益的,例如GEMM采用非平凡的前导维数可以获得更好的性能(参见提示6)。

下图显示了以行主格式保存的大小为rows x columns的二维矩阵的简化情况,其中行具有一些非平凡的(即,不等于列数)步长。
strides

在这种情况下,一般的偏移函数如下所示:

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)。 覆盖所有通道块后,将显示批处理中的下一个图像。

mem_fmt_blk

注意: 我们在格式中使用小写和大写字母来区分块(如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的计算过程中会产生一些额外的计算,但这并不会影响结果。

mem_fmt_padded_blk

所给方法的一些缺陷:

  • 保存数据所需的内存大小不能再通过公式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);

参考资料:

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值