YOLObile:面向移动设备的「实时目标检测」算法

分享一篇AAAI 2021录用论文：YOLObile: Real-Time Object Detection on Mobile Devices via Compression-Compilation Co-Design作者来自于美国东北大学、匹兹堡大学和William & Mary。

作者提出了一种通过从压缩、编译两个角度，在保证模型准确率的基础上，减小模型的大小，并提升模型在移动设备端的运行速度。

通过所提出的YOLObile framework，将YOLOv4压缩了14倍，准确率保持在49.0mAP，在Samsung Galaxy S20上使用GPU，推理速度为17FPS；使用所提出的CPU-GPU合作机制，推理速度可以提升至19FPS，是原始的YOLOv4的5倍。目前代码已经开源。

1. DNN model pruning

这篇文章主要的工作是模型裁剪和推理加速策略，所以在介绍这篇文章的工作之前，先介绍目前主流的三种剪枝策略：Unstructured pruning，Structured pruning和Pattern-based pruning。

1.1 Unstructured pruning

所谓非结构性剪枝允许在权重矩阵的任意位置进行裁剪。其优点主要有:

• 在搜寻最优的剪枝结构上有更好的灵活性
• 可以达到很高的模型压缩率和极低的精度丢失

如下图所示：

可以看出，Unstructured pruning方法得到的权重分布很不规则，所以在计算前向的时候，通常需要额外的非零权值的索引。这对于那些可以并行运算的设备（GPU）很不友好，所以不太适合用于DNN推理加速，甚至有可能导致速度下降。

1.2 Structured pruning

结构性剪枝如上图（b）所示，主要从卷积核个数（filters）和通道数（channels）上进行剪枝，因而得到的权重矩阵仍然是规则的。这对于支持并行运算的硬件非常友好，有助于提升推理速度。

但是由于这种裁剪方式过于粗糙（直接剪掉一个或者多个卷积核或者减去所有卷积核中同一个或多个位置的通道的权重），所以精度丢失非常严重。

1.3 Pattern-based pruning

Pattern-based pruning可以看作是一种fine-grained结构性剪枝，比结构性剪枝更加灵活。如下图所示：

主要包括两个部分：kennel pattern prune和connectivity prune。这里为了方便说明，定义一个卷积层参数为，每个卷积核(kernel)大小为，卷积核个数为。kennel pattern prune指的是对于每一个卷积核，裁去每个通道（channel）上固定位置的参数，如上图黄色背景框内灰色的部分；

每个卷积核可以采用不同的裁剪模式（pattern），但是要保证剩余的每个通道的参数数量是固定的，如上图的剩余的红色、黄色、绿色和紫色方块数量都是4。

而connectivity prune指的是直接裁去整个卷积核。这样做的好处是可以保证裁剪后的参数分布都是规则的，也比structured pruning更加灵活，精度丢失的也会相对少一些。

但是kennel pattern prune只针对3x3卷积核，限制了pattern-based pruning的应用场景。

2. Motivation

基于目前SOTA的目标检测算法，精度高的，模型比较大，在移动设备上会有很高的时延；而那些在移动设备端可以快速运行的轻量级算法又牺牲了算法精度。

三种主流的剪枝算法Unstructured pruning可以保证精度，但是不能保证速度，Structured pruning可以保证速度，但是无法保证精度；Pattern-based pruning可以一定程度上同时保证速度和精度，但是应用场景有限。

基于此，这篇文章的主要工作可以总结为以下两点：

• 提出一种剪枝策略，可以同时保证速度和精度，并且可以推广到任意layer（pattern-based pruning只能应用在3x3卷积层）
• 提出一种更高效的计算加速策略

3. Method

3.1 Block-punched pruning

Block-punched pruning是这篇文章提出的一种以同时保证速度和精度，并且可以推广到任意layer的剪枝策略。主要内容如下图所示：

将一个layer的参数（以卷积层为例）分成参数相等的blocks，每个block包含m个连续的filter和n个连续的channel。对block中的所有channel，裁剪去相同位置的参数。

与pattern-based pruning不同的地方在于，pattern-based pruning对于整个filter进行相同模式的裁剪，而Block-punched pruning对于一个block进行相同模式的裁剪，比pattern-based pruning更加精细，而且可以推广到任意layer（不局限于3x3layer）。

Block-punched pruning特点是：block size会极大的影响模型的精度和在硬件上的运行速度。

• block size越小，精度丢失越少，但是推理速度也会变慢
• block size越大，精度丢失越严重，但是推理速度变快

所以选择一个合理的block size非常重要。这里作者给出了两个建议：

• 对于block中channel的数量：与设备中CPU/GPU的vector registers的长度一致
• 对于block中的filter的数量：在保证目标推理速度的前提下，选择最少的filter数量

3.2 Reweight regularization pruning algorithm

这里讲的是如何具体的实现剪枝。这篇文章采用了一种reweighted group Lasso^[1]^方法，其基本原理是：**减小大权重的惩罚项，增大小权重的惩罚项。**假设是i-th卷积层的参数，于是该问题可以有以下的目标函数来解决：

从公式可以看出，利用参数F范数（Lasso一般采用L1范数）的平方的倒数作为加权值，权值越大，惩罚项的加权值越小。最后需要剪去的参数是那些逼近于0 的参数。

3.3 CPU-GPU合作机制

这里提出了一种更高效的计算加速策略，可以综合利用GPU和CPU。目前的一些推理加速框架如TFLite和MNN只能支持在移动GPU或CPU上顺序执行DNN推理，这可能造成计算资源的浪费。

一些网络如YOLOv4有分不同的分支，如果这些分支可以同时分别运行在GPU和CPU上，就可以提高推理效率和速度。

所以剩下的问题就是确定哪些分支需要执行在GPU上，哪些分支需要执行在CPU上。通常GPU比较适合处理一些高并行的数据，如有大量的卷积层。还是以YOLOv4的cross-stage partial为例子：

定义branch1和branch2在GPU上的耗时为，在CPU上的耗时分别，branch1卷积层多适合GPU运算，如果采用GPU和CPU并行运算，那么最终的处理时间取决于最大耗时，定义数据拷贝到CPU上的耗时为，则GPU和CPU并行运算耗时为：

如果只采用GPU进行串行运算，即先计算branch1，再计算branch2，则耗时为两者之和：

通过和可以确定branch2在哪个设备上运行。因为每个branch的执行是独立的，所以可以通过Greedy Algorithm(贪心算法)来确定网络中每一个分支的执行的位置（GPU or CPU）。

对于那些低计算密度的操作如pixel-wise add和pixel-wise multiply操作，移动设备上CPU和GPU的运算效率差不多。所以对于non-convolution的分支，在CPU还是在GPU上运算，取决于总耗时。

如上图(b)所示，三个YOLO head的运算都是non-convolution的，所以三个分支运算在哪个分支的可能性有8种，假设前两个运行在CPU上，最后一个分支运行在GPU上，那么总的运行时间为:

采用上述的方案分别确定每个conv branch和non-conv branch运行的位置，最小化总的推理时间。

YOLObile提供了每一层的CPU和GPU代码，为实现上述的计算提供了可行性。

4. Results

最终的实验结果表明，文章所提出的压缩剪枝和加速方案可以在提升推理速度的同时，还能保持较高的准确率。

参考文献

[1] EJ Candès, Wakin M B , Boyd S P . Enhancing Sparsity by Reweighted 1 Minimization[J]. Journal of Fourier Analysis & Applications, 2008, 14(5-6):877-905.

[2] 论文原文：https://arxiv.org/abs/2009.05697

[3] 代码：https://github.com/nightsnack/YOLObile