【AutoML--模型搜索】论文阅读：Once-for-All: Train One Network and Specialize it for Efficient Deployment

文章目录

1. 论文概要

论文下载地址：Once-for-All: Train One Network and Specialize it for Efficient Deployment
官方代码：https://github.com/mit-han-lab/once-for-all
发表时间：ICLR2020
作者：MIT 韩松团队

这篇论文的名字可以翻译为《Once-for-All：一次训练，随意部署》，顾名思义，Once-for-All 实现了只训练一个网络，就能够部署于各种硬件平台。不同于以往的 NAS 方法，它实现了模型训练和架构搜索的解耦。

2. 研究背景和动机

深度神经网络（DNN）已经取得了广泛应用，但模型大小和计算成本的爆炸性增长对如何在不同的硬件平台上有效地部署DNN提出了新的挑战，因为它们必须满足不同的硬件效率约束（如延迟和功耗）。

一个常见的例子是，手机上的 App 必须能够支持各种各样的硬件设备，即能够在高中低端的手机上流畅运行。由于不同的硬件资源（如存储器大小、算术单元数），最佳神经网络架构的变化是显著的。即使在相同的硬件上运行，在不同的电池条件或工作负载下，最佳模型架构也有很大的差异。

考虑到不同的硬件平台和部署场景，研究人员必须设计不同的量级的模型来进行部署。即配置网络层的类型（如卷积层、全连接层）、层数、每层的节点数、连接方式等。通常采用的两种方法是：

人为手动设计不同场景下的模型配置。
需要结合不同领域的知识和经验进行设计，并通过多次实验进行比较来获得较好的模型参数。
使用 NAS (neural architecture search) 搜索最佳模型配置。
NAS可以看作是一个优化问题，其目标是在给定的搜索空间中，搜索出最优的神经网络架构，以最大化性能指标（如准确率、速度、模型大小等）并满足资源约束（如计算资源、内存等）。NAS方法通常使用强化学习、进化算法、遗传算法等优化技术来进行搜索。

这两种方式获得的不同模型配置都需要重新训练，这导致为每个场景设计和训练DNN的工程成本和计算代价都是高昂的，并且总成本随着部署场景数量的增加而线性增长。

为了解决上述问题，本文提出了 OFA (Once-for-All)，设计一个 “一次性”“一劳永逸” 的网络，可以直接部署在不同的架构配置下。它灵活地支持不同的深度、宽度、内核大小和分辨率，无需重新训练；并且在给定的部署场景下，选择OFA网络中的特定子网络即可。对于计算资源丰富的云端设备，可以截取一个大的子网；而对于计算资源短缺的边缘设备，则可以截取一个较小的子网。

3. OFA的网络设计

如下图所示，OFA 的总体思想是训练一个包含了很多网络参数配置的大网络 (super-net)，然后评估选出某个场景下最优的子网络 (sub-net)，sub-net 不需要重新训练，只需要 finetune。这相当于 train 和 search 的过程实现了解耦，可以将网络搜索代价由 O(N) 降为 O(1)

假设 OFA 的 super-net 的权重为 $W_O$ ，sub-net 的网络架构配置是 ${arch}_i$ ，可以将网络搜索问题表示为：

$\underset{W_O}{min} \sum_{ {arch}_i} L_{val}(C(W_O,{arch}_i))$

其中 $C(W_O,{arch}_i)$ 表示从 super-net 中选择一部分作为 sub-net 的选择方案。该问题的总体目标是优化 $W_O$ ，使每个 sub-net 与 “独立训练的相同配置的网络” 具有相同的精度。

3.1 Architecture Space (网络配置空间)

super-net 包含了 CNN 中最重要的四种参数：深度、宽度、卷积核大小、分辨率。

根据经验，将 CNN 划分为许多单元 (units)，每个 unit 中的特征图的尺寸逐渐减小而通道数逐渐增加。允许每个 unit 使用任意的深度、宽度、卷积核大小和分辨率。在论文实验中，CNN参数配置为：

输入图像大小范围为 {
    
    min=128, max=224, step=4}
每个unit的可选深度为 {
    
    2,3,4}
每层的的可选宽度膨胀比为 {
    
    3,4,6}
可选的卷积核大小为 {
    
    3,5,7}

因此，对于具有 5 个 unit 的CNN网络，共有 25 种输入分辨率和大约 $\times 3)^2+(3 \times 3)^2+(3 \times 3)^4)^5 \approx 2 \times 10^{19}$ 种网络配置方案。所有 sub-nets 权重共享，否则这么多 sub-nets 将会导致极大的模型参数量。

3.2 训练 OFA 网络

训练 super-net 可以看做一个多目标问题，前者每个优化目标对应一个 sub-net。

3.2.1 Naive Approach

最简单的方法是直接从头开始训练 super-net 网络，通过枚举每个 sub-net 的梯度进行优化。这种方式的训练成本正比于 sub-net 的数量，只适用于有限的 sub-net 的情况，显然不适用于本论文中 $10^{19}$ 量级的网络。

另一种简单的训练方法是在每个更新步骤中采样几个 sub-nets，而不是枚举所有的 sub-nets，这不存在成本过高的问题。但大量的 sub-nets 共享权重，互相干扰，实验发现这会导致明显的精度下降。

3.2.2 Progressive Shrinking (渐进收缩)

super-net 中不同大小的 sub-net 存在嵌套关系，为了防止不同 sub-net 之间的干扰，论文使用渐进的方式执行从大型 sub-net 到小型 sub-net 的训练顺序，称之为渐进收缩。论文中指定了从大网络到小网络的训练顺序：

1、训练sup-net时使用最大的kernel_size = 7，最大depth=4，最大宽度膨胀比=6
2、让kernel_size在{
    
    3,5,7}中选择，深度和宽度依旧保持最大，finetune
3、让depth可以选择，而宽度膨胀比保持最大，finetune
4、让宽度膨胀比也可以选择，finetune

在整个训练过程中，输入图像的分辨率是可选的，通过对每一批训练数据采样不同的图像大小来实现。

渐进收缩与网络剪枝 (Network Pruning) 的区别如下图所示。渐进收缩可以看做一种广义的网络剪枝，它具有更高的灵活性（缩写4种参数而不仅仅是width）。

总结一下，使用渐进收缩的好处有：

可以防止小网络干扰大网络，因此精度更高。因为当 finetune 小网络时，大网络已经收敛了
使用大网络 finetune 小网络可以加快训练流程
渐进收缩提供了较高的模型压缩的灵活度，能够更好地适应不同的硬件平台和效率限制

接下来看一下使用渐进收缩 (PS) 训练 sup-net 的效果。下图展示了八种网络配置在使用和不使用PS时的 ImageNet top1 精度，显然使用 PS 能够获得更好的结果。

3.3 子网络搜索

训练好一个 super-net 后，下一步是根据不同硬件设备的需求和限制搜索出不同的 sub-net。

论文中使用一个孪生神经网络来预测 sub-net 的精度和延迟。具体做法是随机抽样 16K 个具有不同配置和输入图像大小的 sub-nets，在 10K 规模的验证集上计算这些 sub-nets 的精度；然后通过 {architecture, accuracy} pairs 来训练精度预测网络，即预测出给定网络架构和输入图像大小的模型的精度。此外，论文还构建了延迟查找表 (latency lookup table)，以便在给定的硬件平台上预测延迟。

根据精度预测网络和延迟查找表，在给定硬件设置时就可以选择出最优的模型结构。

3.4 附：如何实现CNN的参数压缩

本节介绍论文中如何缩小 [卷积核大小、深度、宽度膨胀比例] 三个参数。

（1）卷积核大小压缩

从大卷积核变为小卷积核只需要对中心的卷积核进行矩阵变换即可，该变换矩阵在 super-net 中参数共享。

（2）深度压缩

如果需要在 N 层的 super-net 中提取出 D层(D<N) 的 sub-net，直接取出前 D 层即可。在 super-net 和 sub-net 中，第一层的权重共享。

（3）通道膨胀比例压缩

下图分别表示 [4,3,2] 个通道的卷积层，根据通道的 $L_1$ 范数衡量每个通道的权重，在通道收缩时保留权重较高的通道。

4. 总结

论文使用与 MobileNetV3 相同的 architecture space 进行实验，使用 ImageNet 进行训练以验证 OFA 的优越性。

OFA的创新点有：

实现了模型训练和架构搜索之间的解耦，在NAS领域具有开创性。
OFA只需单独训练 super-net 而不需要额外训练 sub-net，这极大降低了训练成本。
OFA网络支持CNN中四种参数的变化，为网络设计提供了较高的灵活度。

论文的不足之处在于：

仅通过图像分类任务验证了OFA的有效性，但分类任务是最简单的深度学习应用场景，难以证明OFA在其他更复杂的任务上的有效性。
尽管论文中包含了 $\times 10^{19}$ 个的子网络，在 [128,224] 的输入图像尺寸上进行实验，图像尺寸和参数变化并不大，即很多子网络的差异其实很小。
可解释性的缺失：论文中并没有详细说明选择特定网络结构和渐进收缩过程中参数变化顺序的依据，全文几乎没有理论分析，所以结论均通过实验给出。尽管论文提出的思想具有开创性，但缺乏理论解释且实验的可扩展性不强，因此对于该领域的贡献有待后人证明。