Once-for-All: Train One Network and Specialize it for Efficient Deployment（ICLR2020）

原文链接
源代码

Abstract

我们解决了跨许多设备和资源约束的高效推理的挑战性问题，特别是在边缘设备上。传统方法要么手动设计，要么使用神经架构搜索(NAS)来找到一个专门的神经网络，并针对每种情况从零开始训练它，这在计算上是禁止的(导致二氧化碳排放量多达5辆汽车的使用寿命（哇靠）。在这项工作中，我们建议通过解耦训练和搜索来训练一个支持多种架构设置的一次性(OFA)网络，以降低成本。我们可以快速地从OFA网络中选择一个专门的子网络，而不需要额外的训练。
为了有效地训练OFA网络，我们还提出了一种新的渐进式收缩算法，这是一种广义剪枝方法，它可以在比剪枝(深度、宽度、核大小和分辨率)更多的维度上减小模型大小。它可以获得惊人数量的子网(>10^19)，可以适应不同的硬件平台和延迟限制，同时保持与独立训练相同的精度水平

成果

在不同的边缘设备上，OFA始终优于最先进的(SOTA) NAS方法(与MobileNetV3相比，ImageNet top1精度提高4.0%，或相同的精度，但比MobileNetV3快1.5倍，比EfficientNet w.r.t测量延迟快2.6倍)，同时减少了许多数量级的GPU工作时间和二氧化碳排放。特别是，OFA在移动设置(<600M mac)下实现了新的SOTA 80.0% ImageNet top-1精度。OFA是第三届低功耗计算机视觉挑战赛(LPCVC)、DSP分类赛道和第四届LPCVC(分类赛道和检测赛道)的获奖解决方案。

Introduction

存在的挑战

随着模型尺寸和计算成本的爆炸式增长，如何在不同的硬件平台上有效地部署这些深度学习模型带来了新的挑战，因为它们必须满足不同的硬件效率约束(例如，延迟，能量)

以前的方法

传统方法要么手动设计，要么使用神经架构搜索(NAS)来找到一个专门的神经网络，并针对每种情况从零开始训练它，这在计算上是禁止的(导致二氧化碳排放量多达5辆汽车的使用寿命（哇靠）。这使得它们无法处理大量的硬件设备(231.4亿个物联网设备)和高度动态的部署环境(不同的电池条件，不同的延迟要求等)

本文的方法

本文提出了一种新的解决方案——设计一个一次性的网络OFA，可以直接部署在不同的架构配置下，分摊训练成本。通过只选择一次性网络的一部分来执行推理。它灵活地支持不同的深度、宽度、内核大小和分辨率，而无需重新训练。
图1(左)展示了一个简单的一次性处理(OFA)示例。具体来说，我们解耦了模型训练阶段和神经结构搜索阶段。在模型训练阶段，我们着重于通过选择一次性网络的不同部分来提高所有子网络的准确性。在模型专门化阶段，我们对子网络的子集进行采样，以训练精度预测器和延迟预测器。给定目标硬件和约束，进行预测器引导的架构搜索以获得专门的子网络，成本可以忽略不计。因此，我们将专业神经网络设计的总成本从O(N)降低到O(1)(图1中间)。

在每个更新步骤中枚举所有子网络以获得精确的梯度在计算上是禁止的，而在每个步骤中随机采样几个子网络将导致精度显著下降。挑战在于不同的子网络相互干扰，使得整个一次性网络的训练过程效率低下。为了解决这一挑战，我们提出了一种渐进式收缩算法来训练一次性网络。我们不是直接从零开始优化一次性网络，而是建议首先训练具有最大深度、宽度和内核大小的最大神经网络，然后逐步微调一次性网络以支持与较大网络共享权重的较小子网络。因此，它通过选择较大的子网络中最重要的权重来提供更好的初始化，并有机会提取较小的子网络，从而大大提高了训练效率。从这个角度来看，渐进式收缩可以看作是一种广义的网络修剪方法，它可以收缩整个网络的多个维度(深度、宽度、核大小和分辨率)，而不仅仅是宽度维度。此外，它的目标是保持所有子网的准确性，而不是单一的修剪网络。

我们在许多硬件平台(CPU, GPU, mCPU, mGPU, FPGA加速器)和效率限制下广泛评估了OFA在ImageNet上的有效性。在所有部署场景下，与SOTA硬件感知的NAS方法相比，OFA始终能够显著提高ImageNet的精度，同时节省GPU的时间、成本和二氧化碳排放量。在ImageNet移动设置(小于600M mac)下，OFA在595M mac下实现了新的SOTA 80.0%的top1精度(图2)。据我们所知，这是SOTA ImageNet在移动设置下首次达到80%的top1精度。

Related Work

提出了很多高效神经网络来提高硬件效率
通过神经架构搜索（NAS）来自动化解决准确性和精度的权衡，然而考虑到新的推理硬件平台，这些方法需要重复架构搜索过程并重新训练模型，这导致了令人望而却步的GPU时间、金钱和二氧化碳排放。它们不能扩展到大量的部署场景。单独训练的模型不共享任何权重，导致大的总模型大小和高下载带宽
通过宽度乘法器等实现动态神经网络，相比之下，在这项工作中，我们实现了更多样化的架构空间(深度、宽度、内核大小和分辨率)和更大数量的架构设置。然而，训练网络来实现这种灵活性更具挑战性，这促使我们设计渐进式收缩算法来解决这一挑战

Method

问题描述

假设一次性网络的权值为Wo，体系结构配置为{archi}，那么我们可以将问题形式化为：

其中C(Wo,archi)表示一种选择方案，该方案从一次性网络Wo中选择一部分模型，形成具有体系结构配置archi的子网络。

显然，总的训练目标是优化Wo，使每个支持的子网络保持与独立训练具有相同体系结构配置的网络相同的精度水平。

架构空间

我们的一次性网络只提供一个模型但支持许多不同大小的子网络，涵盖了卷积神经网络(cnn)架构的四个重要维度，即深度、宽度、核大小和分辨率。遵循许多CNN模型的惯例，我们将CNN模型划分为一系列单元，这些单元的特征图大小逐渐减小，通道数逐渐增加。每个单元由一系列层组成，如果特征图大小减小，则只有第一层的步长为2，单元中所有其他层的步幅都是1。

我们允许每个单元使用任意数量的层(表示为弹性深度)；对于每一层，我们允许使用任意数量的通道(表示为弹性宽度)和任意大小的内核(表示为弹性内核大小)。此外，我们还允许CNN模型采用任意输入图像大小(表示为弹性分辨率)。例如，在我们的实验中，输入图像的大小范围为128到224，步幅为4;每个单元的深度从{2,3,4}中选取，每层的宽度展开比从{3,4,6}中选取，核大小从{3,5,7}中选取。因此，使用5个单元，我们大致有**((3 × 3)2 +(3 × 3)3 +(3 × 3)4)5≈2 × 10^19？？？**种不同的神经网络架构，每种结构都可以在25种不同的输入分辨率下使用。由于所有这些子网络共享相同的权重(即，Wo)我们只需要7.7M参数来存储所有参数。如果没有共享，总模型大小将是令人望而却步的

训练OFA网络

Na¨ıve Approach

训练一次性网络可以看作是一个多目标问题，其中每个目标对应一个子网络。从这个角度来看，na¨ıve训练方法是使用总体目标的精确梯度从零开始直接优化一次性网络**，该梯度是通过枚举每个更新步骤中的所有子网络得出的**，如Eq.(1)所示。该方法的成本与子网络的数量成线性关系。因此，它只适用于支持有限数量的子网的场景(Yu et al.， 2019)，而在我们的情况下，采用这种方法在计算上是禁止的

另一种na¨ıve训练方法是在每个更新步骤中抽样几个子网络，而不是枚举所有子网络，这样就不会有成本过高的问题。然而，由于如此多的子网络共享权重，从而相互干扰，我们发现它的精度下降很大

渐进收缩

为了防止子网络之间的干扰，我们建议以渐进的方式从大的子网络到小的子网络强制执行训练顺序，我们将这种培训方案称为渐进式收缩(PS)。
图3和图4提供了PS训练过程的示例，其中我们首先训练具有最大核大小(例如，7)，深度(例如，4)和宽度(例如，6)的最大神经网络。接下来，我们逐步微调网络以支持较小的子网络，逐步将它们添加到采样空间(也可以对较大的子网络进行采样)。即，在训练出最大的网络后（4×6×7），我们首先支持弹性核大小，每层可以从{3,5,7}中选择，而深度和宽度保持最大值。（单一变量法）
然后，我们依次支持弹性深度和弹性宽度，如图3所示。而分辨率在整个训练过程中是弹性的，通过对每批训练数据采样不同的图像大小来实现。我们还在训练最大的神经网络后使用知识蒸馏技术，它使用最大神经网络给出的软标签和真实标签组合两个损失项

与na¨ıve方法相比，PS可以防止小型子网干扰大型子网，因为当一次性网络经过微调以支持小型子网时，大型子网已经得到了良好的训练。对于小子网，它们与大子网共享权重。因此，PS允许用训练良好的大型子网的最重要权重初始化小型子网，从而加快了训练过程。与网络修剪(图4)相比，PS也从训练完整模型开始，但它不仅缩小了宽度维度，还缩小了完整模型的深度、核大小和分辨率维度。此外，PS对大型和小型子网络都进行微调，而不是对单个网络进行微调。因此，与网络修剪相比，PS提供了更强大的一次性网络，可以适应各种硬件平台和效率限制。

PS培训流程详细描述如下：
• 弹性内核大小(图5左)。我们让一个7x7卷积核的中心也可以作为一个5x5核，也可以是一个3x3核。因此，核的大小变得有弹性。挑战在于中心子内核(例如3x3和5x5)是共享的，需要扮演多个角色(独立内核和大内核的一部分)。中心子核的权值可能需要根据不同的角色具有不同的分布或大小。强制它们相同会降低某些子网的性能。因此，我们在共享核权时引入核变换矩阵。我们对不同的层使用单独的核变换矩阵。在每一层中，核变换矩阵在不同的通道之间共享。因此，我们只需要25 × 25 + 9 × 9= 706个额外参数来存储每层的核变换矩阵，这是可以忽略不计的。
• 弹性深度(图5右)。为了在最初具有N层的单元中导出具有D层的子网络，我们保留了前D层并跳过了最后N - D层，而不是像当前NAS方法中那样保留任何D层。因此，一个深度设置只对应于一个图层组合，最后，在大模型和小模型之间共享前D层的权重。
• 弹性宽度(图6)。宽度是指通道的数量。我们赋予每一层选择不同通道扩展比的灵活性。根据渐进收缩方案，我们首先训练一个全宽度模型。然后我们引入一个通道排序操作来支持部分宽度。它根据信道权重的L1范数来计算信道的重要性，对信道进行重新组织。L1范数越大意味着越重要。例如，当从4通道层缩小到3通道层时，我们选择最大的3个通道，其权重与4通道层共享(图6左和中)。因此，较小的子网络被初始化为一次性网络中最重要的通道，该网络已经得到了很好的训练。这种信道排序操作保持了较大子网的准确性。

专业模型部署与一次性网络

训练完一次性的网络后，下一个阶段是为给定的部署场景导出专用的子网。目标是在优化精度的同时，寻找满足目标硬件的效率(例如，延迟，能量)约束的神经网络。由于OFA将模型训练与神经结构搜索解耦，因此在这一阶段我们不需要任何训练成本。此外，在给定神经网络结构的情况下，我们构建了神经网络双胞胎来预测延迟和准确性，为模型质量提供了快速反馈。它通过将测量的精度/延迟替换为预测的精度/延迟(双胞胎)来消除重复的搜索成本。

具体来说，我们随机抽取具有不同架构和输入图像大小的16K个子网络，然后在从原始训练集中抽取的10K验证图像上测量它们的准确性。[结构，精度]是用于训练精度预测器，以预测给定其结构和输入图像尺寸的模型的精度。我们还在每个目标硬件平台上构建了一个延迟查找表来预测延迟。考虑到目标硬件和延迟约束，我们基于神经网络双胞胎进行进化搜索，以获得专门的子网络。

由于神经网络双胞胎的搜索成本可以忽略不计，我们只需要40个GPU小时来收集数据对，并且无论部署场景如何，成本都保持不变

Experiments

在ImageNet上训练OFA网络

PS（渐进式收缩）可以在所有架构设置下显著提高子网络的ImageNet精度

针对不同硬件和约束的专用子网

在处理多个部署场景时，OFA比NAS效率高得多

OFA比NAS方法节省了数量级的设计成本

OFA在ImageNet上取得了最先进的准确性，并且大大降低了延迟

OFA在移动平台上的表现始终优于MobileNetV3，并且注意到OFA可以通过只训练一次(绿色曲线)，在广泛的延迟约束范围内产生具有许多点的整个权衡曲线

专用OFA模型在具有相似延迟的情况下始终能够实现更高的ImageNet精度。更值得注意的是，专门针对一个新的硬件平台并不会增加使用OFA的培训成本

MB4 3x3”表示“扩展比为4，内核大小为3x3的移动块”，它揭示了对不同部署场景使用相同的模型，仅修改宽度乘法器对效率提高的影响有限:随着延迟约束变得严格，准确性会迅速下降

与MobileNetV2和MnasNet相比，OFA模型提高了算法强度(OPS/Byte)和利用率(GOPS/s)

消除内存瓶颈后，利用率和GOPS/s提高了70%-90%

Conclusion

我们提出了一种新的方法OFA，它将模型训练与架构搜索解耦，从而在大量硬件平台下实现高效的深度学习部署。与之前为每个部署场景设计和训练神经网络的方法不同，我们设计了一个一次性的网络，它支持不同的体系结构配置，包括弹性深度、宽度、内核大小和分辨率。与传统方法相比，它降低了训练成本(GPU小时、能耗和二氧化碳排放)。为了防止不同大小的子网络受到干扰，我们提出了一种渐进式收缩算法，使大量的子网络达到与独立训练相同的精度水平。在各种硬件平台和效率限制上的实验证明了我们的方法的有效性。OFA提供了一个自动化的生态系统，可以有效地设计出具有硬件成本模型的高效神经网络