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

1 INTRODUCTION . Deep Neural Networks ( DNNs ) deliver state-of-the-art accuracy in many machine learning applications . However , the explosive growth in model size and computation cost gives rise to new challenges on how to efficiently deploy these deep learning models on diverse hardware platforms , since they have to meet different hardware efficiency constraints ( e.g. , latency , energy ) . For instance , one mobile application on App Store has to support a diverse range of hardware devices , from a high-end Samsung Note10 with a dedicated neural network accelerator to a 5-year-old Samsung S6 with a much slower processor . With different hardware resources ( e.g. , on-chip memory size , # arithmetic units ) , the optimal neural network architecture varies significantly . Even running on the same hardware , under different battery conditions or workloads , the best model architecture also differs a lot . Given different hardware platforms and efficiency constraints ( defined as deployment scenarios ) , researchers either design compact models specialized for mobile ( Howard et al. , 2017 ; Sandler et al. , 2018 ; Zhang et al. , 2018 ) or accelerate the existing models by compression ( Han et al. , 2016 ; He et al. , 2018 ) for efficient deployment . However , designing specialized DNNs for every scenario is engineer-expensive and computationally expensive , either with human-based methods or NAS . Since such methods need to repeat the network design process and retrain the designed network from scratch for each case . Their total cost grows linearly as the number of deployment scenarios increases , which will result in excessive energy consumption and CO2 emission ( Strubell et al. , 2019 ) . It makes them unable to handle the vast amount of hardware devices ( 23.14 billion IoT devices till Published as a conference paper at ICLR 2020 Number of Deployment Scenarios 0 20 40 60 80 16x~1300x reduction direct deploy ( no retrain ) train a once-for-all network specialized sub-nets cpu F P G A Different Hardware / Constraint De sig n Co st Previous : O ( N ) design cost Ours : O ( 1 ) design cost Mobile AI Tiny AI ( AIoT ) Cloud AI MCU 1 20181 ) and highly dynamic deployment environments ( different battery conditions , different latency requirements , etc. ) . This paper introduces a new solution to tackle this challenge – designing a once-for-all network that can be directly deployed under diverse architectural configurations , amortizing the training cost . The inference is performed by selecting only part of the nce-for-all network . It flexibly supports different depths , widths , kernel sizes , and resolutions without retraining . A simple example of Once-for-All ( OFA ) is illustrated in Figure 1 ( left ) . Specifically , we decouple the model training stage and the neural architecture search stage . In the model training stage , we focus on improving the accuracy of all sub-networks that are derived by selecting different parts of the once-for-all network . In the model specialization stage , we sample a subset of sub-networks to train an accuracy predictor and latency predictors . Given the target hardware and constraint , a predictor-guided architecture search ( Liu et al. , 2018 ) is conducted to get a specialized sub-network , and the cost is negligible . As such , we reduce the total cost of specialized neural network design from O ( N ) to O ( 1 ) ( Figure 1 middle ) . However , training the once-for-all network is a non-trivial task , since it requires joint optimization of the weights to maintain the accuracy of a large number of sub-networks ( more than 1019 in our experiments ) . It is computationally prohibitive to enumerate all sub-networks to get the exact gradient in each update step , while randomly sampling a few sub-networks in each step will lead to significant accuracy drops . The challenge is that different sub-networks are interfering with each other , making the training process of the whole once-for-all network inefficient . To address this challenge , we propose a progressive shrinking algorithm for training the once-for-all network . Instead of directly optimizing the once-for-all network from scratch , we propose to first train the largest neural network with maximum depth , width , and kernel size , then progressively fine-tune the once-for-all network to support smaller sub-networks that share weights with the larger ones . As such , it provides better initialization by selecting the most important weights of larger sub-networks , and the opportunity to distill smaller sub-networks , which greatly improves the training efficiency . From this perspective , progressive shrinking can be viewed as a generalized network pruning method that shrinks multiple dimensions ( depth , width , kernel size , and resolution ) of the full network rather than only the width dimension . Besides , it targets on maintaining the accuracy of all sub-networks rather than a single pruned network . We extensively evaluated the effectiveness of OFA on ImageNet with many hardware platforms ( CPU , GPU , mCPU , mGPU , FPGA accelerator ) and efficiency constraints . Under all deployment scenarios , OFA consistently improves the ImageNet accuracy by a significant margin compared to SOTA hardware-aware NAS methods while saving the GPU hours , dollars , and CO2 emission by orders of magnitude . On the ImageNet mobile setting ( less than 600M MACs ) , OFA achieves a new SOTA 80.0 % top1 accuracy with 595M MACs ( Figure 2 ) . To the best of our knowledge , this is the first time that the SOTA ImageNet top1 accuracy reaches 80 % under the mobile setting . 1https : //www.statista.com/statistics/471264/iot-number-of-connected-devices-worldwide/ Published as a conference paper at ICLR 2020 1 2 RELATED WORK . Efficient Deep Learning . Many efficient neural network architectures are proposed to improve the hardware efficiency , such as SqueezeNet ( Iandola et al. , 2016 ) , MobileNets ( Howard et al. , 2017 ; Sandler et al. , 2018 ) , ShuffleNets ( Ma et al. , 2018 ; Zhang et al. , 2018 ) , etc . Orthogonal to architecting efficient neural networks , model compression ( Han et al. , 2016 ) is another very effective technique for efficient deep learning , including network pruning that removes redundant units ( Han et al. , 2015 ) or redundant channels ( He et al. , 2018 ; Liu et al. , 2017 ) , and quantization that reduces the bit width for the weights and activations ( Han et al. , 2016 ; Courbariaux et al. , 2015 ; Zhu et al. , 2017 ) . Neural Architecture Search . Neural architecture search ( NAS ) focuses on automating the architecture design process ( Zoph & Le , 2017 ; Zoph et al. , 2018 ; Real et al. , 2019 ; Cai et al. , 2018a ; Liu et al. , 2019 ) . Early NAS methods ( Zoph et al. , 2018 ; Real et al. , 2019 ; Cai et al. , 2018b ) search for highaccuracy architectures without taking hardware efficiency into consideration . Therefore , the produced architectures ( e.g. , NASNet , AmoebaNet ) are not efficient for inference . Recent hardware-aware NAS methods ( Cai et al. , 2019 ; Tan et al. , 2019 ; Wu et al. , 2019 ) directly incorporate the hardware feedback into architecture search . Hardware-DNN co-design techniques ( Jiang et al. , 2019b ; a ; Hao et al. , 2019 ) jointly optimize neural network architectures and hardware architectures . As a result , they can improve inference efficiency . However , given new inference hardware platforms , these methods need to repeat the architecture search process and retrain the model , leading to prohibitive GPU hours , dollars , and CO2 emission . They are not scalable to a large number of deployment scenarios . The individually trained models do not share any weight , leading to large total model size and high downloading bandwidth . Dynamic Neural Networks . To improve the efficiency of a given neural network , some work explored skipping part of the model based on the input image . For example , Wu et al . ( 2018 ) ; Liu & Deng ( 2018 ) ; Wang et al . ( 2018 ) learn a controller or gating modules to adaptively drop layers ; Huang et al . ( 2018 ) introduce early-exit branches in the computation graph ; Lin et al . ( 2017 ) adaptively prune channels based on the input feature map ; Kuen et al . ( 2018 ) introduce stochastic downsampling point to reduce the feature map size adaptively . Recently , Slimmable Nets ( Yu et al. , 2019 ; Yu & Huang , 2019b ) propose to train a model to support multiple width multipliers ( e.g. , 4 different global width multipliers ) , building upon existing human-designed neural networks ( e.g. , MobileNetV2 0.35 , 0.5 , 0.75 , 1.0 ) . Such methods can adaptively fit different efficiency constraints at runtime , however , still inherit a pre-designed neural network ( e.g. , MobileNet-v2 ) , which limits the degree of flexibility ( e.g. , only width multiplier can adapt ) and the ability in handling new deployment scenarios where the pre-designed neural network is not optimal . In this work , in contrast , we enable a much more diverse architecture space ( depth , width , kernel size , and resolution ) and a significantly larger number of architectural settings ( 1019 v.s . 4 ( Yu et al. , 2019 ) ) . Thanks to the diversity and the large design Published as a conference paper at ICLR 2020 39 train with full width channel importance 0.02 0.15 0.85 0.63 channel sorting progressively shrink the width channel importance 0.82 0.11 0.46 reorg . reorg . progressively shrink the width channel sorting O1 O2 O3 O1 O2 O1 Gradually shrink the width Keep the most important channels when shrinking via channel sorting Progressive Shrinking space , we can derive new specialized neural networks for many different deployment scenarios rather than working on top of an existing neural network that limits the optimization headroom . However , it is more challenging to train the network to achieve this flexibility , which motivates us to design the progressive shrinking algorithm to tackle this challenge . 3 METHOD . 3.1 PROBLEM FORMALIZATION . Assuming the weights of the once-for-all network as Wo and the architectural configurations as { archi } , we then can formalize the problem as min Wo ∑ archi Lval ( C ( Wo , archi ) ) , ( 1 ) where C ( Wo , archi ) denotes a selection scheme that selects part of the model from the once-for-all network Wo to form a sub-network with architectural configuration archi . The overall training objective is to optimize Wo to make each supported sub-network maintain the same level of accuracy as independently training a network with the same architectural configuration . 3.2 ARCHITECTURE SPACE . Our once-for-all network provides one model but supports many sub-networks of different sizes , covering four important dimensions of the convolutional neural networks ( CNNs ) architectures , i.e. , depth , width , kernel size , and resolution . Following the common practice of many CNN models ( He et al. , 2016 ; Sandler et al. , 2018 ; Huang et al. , 2017 ) , we divide a CNN model into a sequence of units with gradually reduced feature map size and increased channel numbers . Each unit consists of a sequence of layers where only the first layer has stride 2 if the feature map size decreases ( Sandler et al. , 2018 ) . All the other layers in the units have stride 1 . We allow each unit to use arbitrary numbers of layers ( denoted as elastic depth ) ; For each layer , we allow to use arbitrary numbers of channels ( denoted as elastic width ) and arbitrary kernel sizes ( denoted as elastic kernel size ) . In addition , we also allow the CNN model to take arbitrary input image sizes ( denoted as elastic resolution ) . For example , in our experiments , the input image size ranges from 128 to 224 with a stride 4 ; the depth of each unit is chosen from { 2 , 3 , 4 } ; the width expansion ratio in each layer is chosen from { 3 , 4 , 6 } ; the kernel size is chosen from { 3 , 5 , 7 } . Therefore , with 5 units , we have roughly ( ( 3× 3 ) 2 + ( 3× 3 ) 3 + ( 3× 3 ) 4 ) 5 ≈ 2× 1019 different neural network architectures and each of them can be used under 25 different input resolutions . Since all of these sub-networks share the same weights ( i.e. , Wo ) ( Cheung et al. , 2019 ) , we only require 7.7M parameters to store all of them . Without sharing , the total model size will be prohibitive .

Summary: In this papers, the authors learn a Once-for-all net. This starts as a big neural network which is trained normally (albeit with input images of different resolutions). It is then fine-tuned while sampling sub-networks with progressively smaller kernels, then lower depth, then width (while still sampling larger networks occasionally, as it reads). This results in a network from which one can extract sub-networks for various resource constraints (latency, memory etc.) that perform well without a need for retraining.