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

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

深度学习模型变得越来越大，越来越复杂，由于其有限的计算和能源预算，部署在低功耗电话和IoT设备上变得充满挑战。深度学习的最新趋势是使用高度量化的模型，该模型可对输入和几位权重进行操作，诸如XNOR-Net，DoReFa-Net和HWGQ-Net之类的网络正在稳步发展，提高了准确性。

下面是一个低精度图形片段的示例。低精度卷积将量化的数据和位包放入适当的数据布局中，实现有效的比特卷积。输出具有更高的精度，在对其进行重新量化，另一个低精度算子发送之前，将传统的深度学习层（如批处理归一化和ReLu）应用于该输出。

 

 低精度卷积管线。

理论上讲，低精度算子比浮点算子使用更少的运算，可以实现巨大的加速。深度学习框架难以置信的良好优化的低级BLAS和LAPACK库，利用了数十年的工程工作，CPU包含用于加速这些任务的内在指令。开发与8位量化甚至浮点算子的卷积之类的低级算子并不简单。本文介绍了为CPU自动生成优化的低精度卷积的方法。声明低精度算子，根据有效存储的低精度输入进行计算，描述用于描述实现参数搜索空间的调度。依靠AutoTVM来快速搜索空间，找到针对特定卷积，精度和后端的优化参数。

二进制计算背景

低精度模型的核心是biterial点积，使得仅使用按位运算和popcount即可计算卷积和密集算子。点积是两个向量的元素逐次相乘，将所有元素相加来计算的，就像下面的简单示例一样。如果所有数据都是二进制数据，可以将输入向量打包为单个整数，按位与打包的输入，使用popcount对结果中的1进行计数来计算点积。根据输入数据的量化方式，可以使用按位xnor代替按位二进制点积。

 

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

 

Bitserial点积

在TVM中定义算子

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

 

 不同的位压缩布局

一旦采用这种位打包格式，就可以按位计算低精度卷积。该数据沿输入通道打包，位平面被添加到最里面的轴，数据打包为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的张量化原语调度。由于编译器会产生次优的代码，人们经常可以编写较短的汇编序列，提高效率。这些微内核利用引入的新内在函数，帮助加速深度学习工作负载，使用巧妙的方式来改善内存访问或减少所需的指令数量。

结果

树莓派Raspberry

与16位整数TVM实现相比，Raspberry Pi 3B的卷积速度提高了。工作负载是ResNet18的卷积层。

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

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

手动优化实现2位权重1位激活Raspberry Pi卷积的速度。

x86

与32位浮点TVM实现相比，x86上的卷积速度有所提高。x86不支持此微体系结构的矢量化popcount，加速性较低。

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