12倍端到端加速,陈天奇创业公司OctoML提出克服二值网络瓶颈新方法

Riptide 是一种新的模型量化方法,可以将模型量化至 1、2 位。研究团队今年三月在 MLSys 上介绍了 Riptide,这篇文章主要讲一下为什么要构建 Riptide,并快速了解它的幕后工作原理。团队计划来年将 Automatic ultra low-bit 功能添加到 Octomizer 中。在此之前,读者可以使用开源 Riptide 项目和 MLSys 论文中的信息来进行模型优化。

  • 论文链接:https://proceedings.mlsys.org/static/paper_files/mlsys/2020/155-Paper.pdf

  • GitHub 项目:https://github.com/jwfromm/Riptide

动机及背景

机器学习发展迅速。几乎每个月,那些优秀的新模型都会大大提高一些视觉或语言任务方面的 SOTA。这些改进一部分是由新的算法和架构创新所推动的,但对于深度学习任务来说,不断扩展的算力和内存也使其取得了重大进展。

随着 ML 准确率的提升,模型所需要的算力和内存也不断增加。

早在 2016 年,我们就可以看出模型大小和准确率之间的关系了。当下,许多 SOTA 模型只能在最新的 NVIDIA GPU(或 GPU 集群)上才能有效运行。而许多用户无法花费数千美元在高端 GPU 上,所以就限制了模型部署的范围。

由于各类网络可能只需要训练一次,开发人员也许可以证明这种训练费用是合理的、有效的。但是,一旦模型进行了部署,将长时间被大量用户不断运行。要以这种规模部署最新的模型,通常需要将输入数据流传输到云上,并将预测返回到用户设备。现在许多应用程序都依赖这种方法,但它有一些缺点,体现在网络连接性、延迟、隐私问题以及庞大复杂的基础架构上。

为了避免这些缺点,许多团队探索如何在低成本的终端用户硬件(如手机或 IoT 设备)上直接运行最新模型。这是一个巨大的挑战,因为此类设备没有足够的算力或内存。例如,树莓派 3 比 NVIDIA Titan GPU 慢了大约 4000 倍。

为了使高精度模型适应此类平台,最近的研究方向已经开始探索如何使这类网络运行更快,同时占用更少的内存。从较高的层面来说,这些技术遵循两种策略:体系架构优化和近似优化。架构优化涉及寻找连接层的新方法,以减少延迟或提高参数有效性。MobileNet 和 SqueezeNet 是两个以移动端为重点的体系架构。与创建新的移动设备友好型模型相反,近似优化旨通过加快运算速度提高现有模型的速度,同时保持足够的准确率。

两种流行的近似优化方法。

在最近流行的几种近似优化中,知识蒸馏和剪枝是两种有代表性的方法(如上图所示)。前者尝试使用大型教师网络来更好地训练学生网络;后者则删除了网络中影响较小的权重和激活函数。本文将重点介绍另外一种方法——二值神经网络,这种网络在泛化性能、加速潜力、内存压缩等方面都有优秀的表现。

二值网络

为了提高性能并减少内存需求,研究团队在部署模型时越来越多地量化激活函数和权重。例如,在推理过程中,工程师可能将模型转换为能够对 int8(8 位整数)进行运算的方式,而不是在训练过程中经常使用的 float32(IEEE 754 单精度浮点数)。小型整数运算不仅比浮点运算速度快,更因为它们占用的位数更少,所以可以通过充分利用可用内存带宽来提高吞吐量。在实践中,很多模型被量化后都不会有显著的准确度损失,因此这项技术大受欢迎。

二值化(binarization)将量化进行到了极致,将网络的权重和激活函数降低到只剩一位,这会带来一些新的优化。考虑两个 1 位值之间所有可能的乘法(如下图),这类方程式和「与门」的逻辑真值表极为相似。

如果认为值为 0 的位代表-1,则上表将成为「同或门」的真值表。

这种等效可以让我们用更高效的二值运算代替浮点运算。

浮点点积(上)和二值点积(下)的比较。

对上图点积内循环中的运算次数进行统计可以发现,二值化可以将一层中的运算次数减少到原来的 1/43,将参数大小减少至 1/32。当这样大体量的优化效果首次由 Courbarioux 等人提出时,2016 年便掀起了二值网络研究的热潮。当然,用 1 位值近似 32 位浮点数是一种有损近似。与对应的全精度相比,二值网络通常会有明显的精度损失,top-1 准确率通常会损失近 20%。因此,二值网络的研究重点一直聚焦于如何减少精度损失。

尽管多个研究团队在提高二值网络准确率方面取得了巨大进步,但他们都没有以一种可以衡量端到端加速的方式实现该网络。这种实现的缺乏不仅让我们很难知道二值网络的实际速度,而且还阻止了二值网络在许多实际环境中的应用。

而缺乏实现的原因是,要编写一个简单的二值矩阵乘法,尽管运算次数比较少,但它依然会比大多数的全精度乘法慢得多。函数的运行时间并不完全取决于其运行次数,内存的访问方式也起着重要作用。任何二值网络的实现都必须与 OpenBLAS 和 MKLDNN 这样的库竞争,而这些库经历了大型工程团队多年的手工优化,难度可行而知。而对于大多数研究机构来说,花大量的时间和精力来建立具有竞争力的库根本不可能。取而代之的是,他们在训练期间模拟二值化,并假设根据运算次数可以预测加速。

Riptide 的突围

为了解决这些问题, OctoML 的研究者提出了 Riptide,这是一种找出并解决端到端二值网络瓶颈的 方法。Riptide 基于 TVM,后者是一种深度学习系统编译器,可以帮助我们自动生成经过调优的高性能二值化算子

迄今为止,二值网络优化仅仅着眼于高效实现核心低位(low-bit)卷积层的不同策略。这一着眼点基于以下假设,即二值网络性能可以反映高精度网络的行为:如果核心卷积可以利用尽可能少的位数实现足够高的准确率,那么整个网络将变得非常快。然而,没有哪个二值网络是单纯地由卷积构成。在卷积之间有许多关键的中间运算,需要用它们来处理下一层的数据。在高精度网络中,这些层的延迟可以忽略不计,但在二值网络中,卷积可以实现 43 倍的加速,这些中间「粘合层(glue layers)」就变得非常重要。

二值卷积之间的「粘合层」及其计算复杂度。H 和 W 表示输入维度,F 表示滤波器的数量。

当前多数的二值网络至少包含 4 个上图中的蓝色层(QConv 和 Quantize 之间)。从左到右来看,首先将 QConv 的输出从整数形式去量化为等效的浮点数,接下来,使用权重缩放(传播实值权重的大小)和批归一化(保持激活的分布可预测)进行缩放。然后应用 ReLU 等非线性激活函数,结果被重新量化为单个位,并打包为下一个量化卷积做准备。

为了理解粘合层的重要性,可以联想一下 SqueezeNet(一种更高效的移动端部署架构)。假设典型的输入大小约为 200×200 像素,并且二值卷积的完美 roofline 实现运行速度是全精度的 43 倍,那么可以估计出网络在粘合层中花费的总执行时间。

假设二值化可以使卷积的速度提高近 43 倍,可以进一步估计,更高精度的粘合层将消耗大约 70% 的总推理时间。这是一个相当大的瓶颈!即使在较低的卷积速度(如 20 倍和 10 倍),粘合层仍然会消耗大约一半的推理时间。因此,研究者认为,要真正实现二值网络所承诺的加速,粘合层也必须是二值化的!

前面已经介绍过融合粘合运算(fused glue operation),它只用两个指令就完全替代了粘合层。细节可以参见论文。其关键思想是用移位运算代替乘法,将缩放项近似为 2 的近似幂,用定点量化近似(fixed point quantized approximations)代替浮点加法和减法。定点量化近似可以直接添加到二值卷积输出中。综上所述,得到如下方程组:

其中,N 是用来量化网络激活的位数,最后一行求解 q(a) 给出了融合粘合的完整方程。通过替换这种融合粘合运算,可以创建一个完全二值化的网络:

研究者还在 ImageNet 数据集上进行了广泛的准确度扫描,发现与其他 SOTA 二值网络相比,上文中的融合粘合运算不会造成任何准确度损失。

这是一个好消息,因为在体系架构层次上消除了粘合瓶颈之后,现在已经可以着手生成快速的二值卷积并度量端到端性能。

端到端加速

早期,研究者将二值网络作为一种近似技术,帮助我们在移动和 IoT 设备上部署高效的模型。先来看一下 Riptide 在树莓派 3b 上的表现如何。树莓派 3b 基于 ARM Cortex-A53 处理器,是资源受限环境的二值目标的典型代表。

首先,将基于 for 循环的简单网络实现与使用 MKLDNN 来加速其浮点运算的完整精度基线进行比较。

使用 MKLDNN 的完整精度 ResNet18 运行时间 vs. 未优化的完全二值网络。

在没有优化的情况下,研究者获得了 ResNet18 的一个完全二值网络(FBN)版本,其速度大约相当于对应高精度网络的 1/5。为了真正利用二值网络,现在还需要生成一个经过调优的执行 schedule。研究者通过在 TVM 中实现二值算子并利用其强大的调度、优化能力来得到这种 schedule。

TVM 提供了以下可以用于优化函数的 schedule 特性:

  • Tiling 将一个计算分解成多个块,以改善负载的内存局部性。

  • Vectorization 利用硬件 SIMD 指令来实现更高效的运算执行。

  • Parallelization 利用多核等 MIMD 设施。

  • Loop Unrolling 复制循环体来减少开销。

此外,研究者还利用了《Automating Generation of Low Precision Deep Learning Operators》论文中提到的 fast popcount 算子。最后,他们还介绍了一种名为 bitpack fusion 的新优化方法。

Bitpack fusion 将 bitpacking 折叠成卷积,以减少激活内存消耗。

在高级层次上,bitpack fusion 尽可能将 bitpacking 折叠到前面的卷积核中。这可以将中间内存需求减少到原来的 1/16,使得复制策略更加高效,从而提供进一步的加速。

那么,这一堆优化究竟哪些更重要呢?

为了解答这一问题,研究者在 SqueezeNet 上进行了控制变量研究,结果如上图所示。

他们发现,与更高精度的浮点基线相比,每种优化都对加速有显著影响。如果将各种优化方法一起应用,可以在真实模型上看到 10 倍的加速。虽然 10 倍的加速远小于论文中所期望的 43 倍加速,但这也是一个数量级的提升,可以显著扩展能够高效运行 SOTA 模型的设备范围。当然,这种加速只有在保证准确率的前提下才有价值。

上表显示了几种流行二值模型在 1、2、3 位激活条件下的准确率与运行时间权衡。总体来看,可以发现在提供更多加速的同时,Riptide FBN 和其他 SOTA 二值技术一样稳定,甚至比后者更加准确。Riptide FBN 可以提供准确的激活位宽和量化极性(quantization polarity),提供 4~12 倍的加速,使得它很容易满足所有应用的准确性需求。

原文链接:https://medium.com/octoml/riptide-fast-full-binarization-in-tvm-ae2afd2104bb

文为机器之心编译, 转载请联系本公众号获得授权

✄————————————————

加入机器之心(全职记者 / 实习生): hr@jiqizhixin.com

投稿或寻求报道:content @jiqizhixin.com