[NAS]Once-For-All:Train One Network and Specialize itfor Efficient Deployment

from MIT HAN Lab

ICLR2020

Abstract

我们已经在跨设备或资源受限的边缘设备上完成了搞笑推理。传统的方法亦或人工设计亦或通过NAS搜索神经网络并从头开始训练使用不同任务的网络，但搜索成本十分高昂，比如CO2的排放相当于五辆车，因此不易于扩展和应用。本文提出一种一次性训练网络OFA，将网络的搜索和训练解耦开来支持不同的架构设计从而降低成本。本文通过OFA网络中快速的选择子网络来实现，并提出渐进式收缩算法，这是一种广义的剪枝算法，比常规剪枝在更多维度（如深度、宽度、卷积核大小、分辨率）方面都可以减少模型大小，同时能获得数量惊人的子网络（>10^19）从而适应不同硬件平台和延迟约束，并且保证与独立训练网络相同的精度水平。
在不同的边缘设备上OFA网路的性能使用优于SOTA的NAS算法，比如比MobileNetV3精度高4%，相同精度下比MobileNetV3推理速度快1.5x。OFA在Imagenet达到了80.0%的分类精度。

Section I Introduction

DNN在应用于诸多机器学习应用中，但是模型规模和计算成本的爆炸增长给如何在不同硬件平台上部署提出了挑战，因为他们必须要满足不同的硬件效率约束，如延迟、功耗。对于不同的硬件资源，如计算单元、芯片内存大小最优的网络结构往往差别较大。即使运行在相同硬件上在不同的电池条件或工作负载下的表现也是不同的。
鉴于不同的硬件凭条和效率限制（称之为部署场景）研究人员设计了不同的紧凑模型，有的通过模型压缩等实现高效部署。但是为每一个应用场景都进行定制化设计人工成本十分高昂，而且有的方法基于重复的网络设计古城，需要在每种情况下对设计的网络从头训练，总的设计成本会随着部署场景的增加而线性增长，会导致过度的能源消耗和温室气体的排放。

本文提出一种新的解决方案，设计一个一次性的网络然后部署在不同的硬件平台，从而分摊设计成本。推理时只选用母网络的一部分，并且可以灵活配置网络的深度、宽度、卷积核大小和分辨率而无需在训练。Fig1 左侧展示了一个OFA的实例，主要是解耦了网路的搜索和训练阶段，在模型训练阶段关注设计OFA网络解雇从而提升所有子网络的性能；在模型定制化阶段则关注采样子网络的子集训练精度预测器和延迟预测器。从而在目标硬件和约束条件下获得符合约束的子网络，成本可以忽略不计，将设计总成本从O（N）下降到O(1).

但是OFA的训练并不简单，因为需要联合优化权重从而保证大量子网络的性能（超过10^19）.在计算过程中枚举所有子网络计算梯度是不显示的，而每一步随机采样子网络更新会带来精度的下降，因为不同的子网络之间会相互干扰。

为了解决这一挑战本文使用渐进的收缩算法来训练OFA而不是从头开始直接优化OFA。而是首先训练具有最大深度、宽度和内核大小的最大的神经网络，然后逐渐对子网络进行优化，从而支持与更大规模的子网络共享一部分权重，这样可以显著提升训练效率。从这一角度出发，渐进收缩可以看做是剪枝的一种，但是会缩减网络的多个维度，不仅仅是宽度，此外收缩的目标是全体子网络的性能而不是单一网络的性能。
本文在诸多硬件平台，如GPU,CPU,mCPU,mGPU,FPGA加速器上进行了测试，在所有部署场景下与SOTA硬件感知的NAS方法相比OFA获得了更高的精度，并且首次在移动端ImageNet获得了超过80%的准确率。

Section II Related Work

Efficient Deep Learning 


高效率网络：SqueezeNet.MobileNet,ShuffleNet


网络效率提升：模型压缩，如剪枝去除冗余的神经元或通道，量化减少权重和激活值的位宽


NAS


目前有越来越堵的NAS算法会将硬件效率考虑在内。但是对于新的硬件平台有需要重复这一搜索过程增加了搜索成本，因为每次训练的模型不共享任何权重导致模型总规模较大。


Dynamic Neural Networks



在动态神经网路哦方面有的使用控制器有的在计算图中引入早期退出分支，还有基于特征图自适应修剪通道或者引入随机的降采样点来减少特征图大小。本文在一个更多样的网路结构空间中从而为不同的部署场景退出专用神经网路，而不是基于现有的神经网络进行缩放。

Section III Method

Part 1 Problem formalization

Wo表示OFA网路哦，C表示采样策略，负责从Wo中生成子网，那么问题就是求解每一个子网络与独立训练的子网络保持相同的精度。

Part 2 Architecture Space

OFA只有一个网路模型但是具备不同规格的子网络，主要在深度、宽度、卷积核维度和分辨率上有所不同。OFA的网络结构特征图会逐渐减少，通道数逐渐增加，每一个unit中包含一系列层，只有第一层步长为2用于下采样。
其中网络宽度、深度、卷积核大小都是多种的
输入图像

分辨率从128到224 4种步长


深度：2 3 4


宽度：3 4 6 



卷积核 3 5 7



因此可能的子网络有：

Part 3 Trainging the Once-For-All Network

Naive Approach 




OFA的训练可以看做是一个多目标任务，每一个子网络就是一个目标，因此原始的训练方法就是逐一优化子网络，参见Eq1，那么搜索成本也就会随着子网络的数目线性增长。
另一种训练方法是每次更新时只采样一部分子网络进行更新，但是由于子网络之间权重共享因此会互相影响，本文提出progressive shrinking的方式来进行优化。

Progressive shrinking

OFA网路欧中包含许多不同大小的子网路，小的子网路嵌套在大的子网路中，为了防止子网之间的干扰本文建议从大的子网训练到小的子网，本文称这种训练方式为渐进收缩（PS），参见Fig 3-4，可以看到会最开始从最大卷积核（7）、最大宽度（6）和最大深度（4）开始训练，然后逐步微调网路偶来获得更小的子网络进入采样空间，此时更大的子网络也会被采样）。

具体来说在训练完最大的子网后，会先从卷积核中选择，此时网络宽度和深度保持最大值；随后在进行网路哦宽度和深度的选择，整个搜索过程是一个弹性的，分辨率也是弹性的通过每批次训练数据resize实现。
与原始训练方法相比，PS可以防止小的子网影响大的子网，因为大的子网已经得到了良好的训练，PS会将训练好的大的子网继承给小的子网，这也加速了训练古城，与Fig 4的网络修建相比，PS从完整的网络模型开始，在缩小宽度维度时拥有完整的网路哦深度、卷积核和分辨率。并且PS是基于微调而不是使用剪枝后的网路，这样PS最终得到的结果是一个性能强大的OFA网络，可以根据不同的硬件平台和约束条件进行微调。

PS流程表述如下：

Elastic Kernel Size

Fig 5展示了弹性卷积核，可以看到7x7卷积核的中心是5x5卷积，5x5卷积中心又是3x3卷积，因此卷积的中心部分具有不同的角色，确实强迫这部分的一致性会降低一些子网络的性能，因此本文引入共享权值部分的转换矩阵，通过这一转换矩阵获取不同层对应的权重，这样只需要25x25+9x9个额外的参数来储存每一层的卷积核变换矩阵。

Elastic Depth


如果要从一个N层网路中获取D层，本文会保持前D层然后跳过后面的N-D层，而不是像其他NAS算法中随机保留D层。这样D层的权重在大小子网之间是共享的。



Elastic Width 




网络宽度指的就是通道数。在PS中本文先训练全宽度的最大子网，然后引入一个通道排序操作，会根据通道权重选择某些通道。L1范数越大意味着当前通道权重更高，会被选择，这样小的子网可以很好的保留较大子网的权重，保证了小的子网的性能。


Part 4 Specialized model deployment with OFA Network

OFA训练好后接下来的任务就是推导出用于给定部署场景的子网，目标往往是满足某些指标，如效率、延迟、功耗的神经网络。本文使用NN-Twins来预测延迟和精度，这样将重复的延迟/精度预测转换为通过predictor实现。
具体来说，本文弧随机抽取16K个不同结构不同输入分辨率的子网，然后在10K个验证图像上测试精度，从而获得[结构，精度]对用于训练精读预测期；延迟的数据则是建立一个查找表实现，然后基于NN-twins进行进化搜索来获得一个特定的子网，而使用NN-twins的搜索成本可以忽略不计。

Section IV Experiments

本文首先在ImageNet上使用PS训练OFA，随后将训练好的OFA在不同硬件平台上进行测试，包括：Samsung S7 Edge, Note8, Note10, Google Pixel1, Pixel2,
LG G8, NVIDIA 1080Ti, V100 GPUs, Jetson TX2, Intel Xeon CPU, Xilinx ZU9EG, and ZU3EG
FPGA

Part 1 Training the OFA network on ImageNet

OFA全网络 32GPU 180 epochs batch_size 2048
训练时间 1200 GPU Hours V100


Results 


Fig 7展示了PS搜索到的OFA网络其中8种子网络的结果，可以看到PS对所有规模上的精度都有提升，比如使用PS 在MobileNetV3-Large级别达到了74.8%的top-1 Acc，而不使用PS精度只有71.5%



Part 2 Specialized Sub-networks for different hardware and constraints

Table 1展示了OFA与当前NAS算法在移动端上的性能对比，可以看到OFA更高效，计算成本基本不变，而其他的NAS算法会与部署场景线性增长。比如N=40时OFA的碳排放量比ProxylessNAS减少了16倍，比FBNet减少了19倍，比MNASNet减少了1300倍。
不适用微调PS原始精度能到76%，进行微调25个epoch精度会提升至76.4%；使用PS比不适用PS精度提升3.6%
Fig 9展示了不同计算资源约束下的OFA性能对比，可以看到389M MACs时OFA的精度是79.1%，595M下精度为80.0%，比EfficientB2降低了1.68x。尤其是OFA在硬件上臂EfficientNet快得多，提升了2.6倍。我们还发现从头训练得到的神经网络性能也比不过OFA，这说明不仅仅是网络结构，网络训练的权重也有助于提升OFA的性能。
Fig 10展示了OFA与MobileNet在移动端的详细对比，可以看到OFA可以在很大范围内进行调控，只用训练一次，这在之前的NAS是不可能的，因为徐亚高昂的训练成本。

OFA for diverse Hardware platforms

Fig 11展示了不同硬件平台上的对比，OFA也超过了MobileNet。可以很好的权衡延迟和精度，Fig 14展示了搜索到的网络具体结构。可以看到仅仅修改宽度multiplier对效率的提升有限，随着演出约束越来越高模型精度会迅速下降。

OFA for specialized Hardware accelerators


很少有NAS测试在硬件加速器上的性能，本文在两种FPGA上进行了测试。
首先评价的是计算强度（OPs/Bytes），OPs/Bytes越高，内存瓶颈越小，越容易并行化，参见Fig 12，OFA比MobileNet高48%，使得资源利用率更高。
CNN的网络结构也应该与硬件加速器进行协同设计，Fig 14显示在FPGA只对3x3有较好的支持，大核卷积会导致高昂的计算成本；在CPU上则超过50%的计算都是用了大核卷积；而FPGA和GPU因为并行化比CPU好，所以设计出的模型也比CPU的更宽。

Section V Conclusion

本文提出的OFA算法将网络搜索与网络训练解耦，可以在不通过硬件平台上实现高效的部署，通过只训练一次，支持弹性深度、宽度、卷积核、分辨率有效降低了训练成本；本文还提出了渐进搜索策略使得大量子网络可以达到与单独训练相同的精度水平，OFA提供了一种自动化的利用硬件资源的高效设计模型。