自动化生成低精度深度学习算子

•


随着深度学习模型变得越来越大和复杂，由于其有限的计算和能源预算，将它们部署在低功耗手机和物联网设备上变得具有挑战性。深度学习领域最近的一个趋势是使用极度量化的模型，这些模型在几位的输入和权重上运行，像 XNOR-Net、DoReFa-Net 和 HWGQ-Net 这样的网络在提高准确性方面取得了稳步进展。

下面是一个低精度图代码片段的示例。低精度卷积接收量化数据并将其位打包成适当的数据布局，以实现高效的位串行卷积。输出具有更高的精度，并且在其被重新量化并通过另一个低精度算子之前，对其应用了传统的深度学习层，例如批归一化和 ReLu。

image

低精度卷积流水线。

理论上，低精度算子比浮点算子使用更少的操作，这使得许多人相信它们可以实现巨大的加速。然而，深度学习框架利用了数十年通过低级 BLAS 和 LAPACK 库的工程工作，这些库经过了极其良好的优化，并且 CPU 包含加速这些任务的内在指令。在实践中，开发诸如卷积之类的低级算子以与 8 位量化甚至浮点算子竞争并非易事。在这篇文章中，我们介绍了我们自动为 CPU 生成优化的低精度卷积的方法。我们声明我们的低精度算子，以便它们在有效存储的低精度输入上进行计算，并描述一个调度，该调度描述了实现参数的搜索空间。我们依靠 AutoTVM 快速搜索空间并为特定的卷积、精度和后端找到优化的参数。

位串行计算背景

低精度模型的核心是位串行点积，它使卷积和密集算子能够仅使用按位运算和 popcount 进行计算。通常，点积是通过两个向量的元素wise乘法，然后对所有元素求和来计算的，如下面的简单示例所示。如果所有数据都是二进制的，则可以将输入向量打包成单个整数，并且可以通过对打包的输入进行按位与运算，并使用 popcount 计算结果中 1 的数量来计算点积。注意：根据输入数据的量化方式，可以使用按位异或代替按位与。

image

二进制点积。

任意精度的点积可以通过首先将输入数据分离成位平面以这种方式计算。一旦以这种表示形式，我们可以通过对 A 和 B 的位平面之间的加权二进制点积求和来计算点积。二进制点积的数量随着 A 和 B 精度的乘积而增长，因此这种方法仅适用于非常低精度的数据。

image

位串行点积。

在 TVM 中定义算子

在计算之前，需要对输入数据进行位打包，以便可以访问输入数据的位平面，并将其打包成受支持的数据类型，例如 uint8 或 uint32。我们提供了一个灵活的位打包算子，它可以接受任意大小的输入张量，并返回一个位打包张量，用户可以在其中指定位平面应位于哪个轴上。

image

不同的位打包布局。

一旦采用这种位打包格式，就可以按位串行方式计算低精度卷积。对于此演示，数据沿着输入通道打包，位平面被添加到最内层轴，并且数据被打包成 32 位整数。位串行卷积的计算方式类似于普通卷积，但是按位与 (&) 替换了乘法，并且我们使用 popcount 来累积打包数据中的值。位平面轴成为额外的归约轴，并计算输入和内核的不同位平面之间的二进制点积。最后，输出以解包格式和更高的精度计算。

Input_bitpacked = bitpack(Input, activation_bits, pack_axis=3, bit_axis=4, pack_type=’uint32’)
Weights_bitpacked = bitpack(Filter, weight_bits, pack_axis=2, bit_axis=4, pack_type=’uint32’)
batch, in_height, in_width, in_channel_q, _ = Input_bitpacked.shape
kernel_h, kernel_w, _, num_filter, _ = Filter_bitpakced.shape

stride_h, stride_w = stride
pad_top, pad_left, pad_down, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))

# Computing the output shape
out_channel = num_filter
out_height = simplify((in_height - kernel_h + pad_top + pad_down) // stride_h + 1)
out_width = simplify((in_width - kernel_w + pad_left + pad_right) // stride_w + 1)
pad_before = [0, pad_top, pad_left, 0, 0]
pad_after = [0, pad_down, pad_right, 0, 0]
Input_padded = pad(Input_bitpacked, pad_before, pad_after, name="PaddedInput")

# Treat the bitplane axes like additional reduction axes
rc = tvm.reduce_axis((0, in_channel_q), name='rc')
ry = tvm.reduce_axis((0, kernel_h), name='ry')
rx = tvm.reduce_axis((0, kernel_w), name='rx')
ib = tvm.reduce_axis((0, input_bits), name='ib')
wb = tvm.reduce_axis((0, weight_bits), name='wb')


tvm.compute((batch, out_height, out_width, out_channel), lambda nn, yy, xx, ff:
             tvm.sum(tvm.popcount(
               Input_padded[nn, yy * stride_h + ry, xx * stride_w + rx, rc, ib] &
               Weights_bitpacked[ry, rx, rc, ff, wb])) << (ib+wb))).astype(out_dtype),
               axis=[rc, ry, rx, wb, ib]))

在我们的调度中，我们应用了常见的优化，例如向量化和内存平铺，以提供更好的内存局部性并利用 SIMD 单元。 诸如平铺之类的某些优化需要针对特定微架构进行调整的参数。我们将这些参数作为旋钮公开给 TVM，并使用 AutoTVM 自动同时调整所有参数。

最后，我们可以制作小的微内核来替换计算的最内层循环，并使用 TVM 的 tensorize 原语来调度它们。由于编译器通常会生成次优代码，因此人们通常可以编写更高效的短汇编序列。这些微内核通常利用正在引入的新内部函数来帮助加速深度学习工作负载，并巧妙地使用它们来改善内存访问或减少所需的指令数量。

结果

Raspberry Pi

与 16 位整数 TVM 实现相比，Raspberry Pi 3B 上的卷积加速。工作负载来自 ResNet18 的卷积层。

image

与 16 位 TVM 实现相比，Raspberry Pi 上低精度卷积的加速。

与来自 移动设备上的高性能超低精度卷积 的手动优化实现相比，Raspberry Pi 3B 上的 2 位激活、1 位权重卷积加速。工作负载来自 ResNet18 的卷积层。

image

针对手动优化实现，Raspberry Pi 上 2 位权重 1 位激活卷积的加速。

x86

与 32 位浮点 TVM 实现相比，x86 上的卷积加速。注意：x86 不支持此微架构的向量化 popcount，因此加速较低。

image

与 32 位浮点 TVM 实现相比，x86 低精度卷积的加速。

给我展示代码

参考文献