Fast and Efficient Once-For-All Networks for Diverse Hardware Deployment

1 INTRODUCTION . Convolutional neural networks ( CNNs ) are overwhelmingly successful in many machine learning applications . These applications may have different inference constraints ( e.g. , latency ) and are deployed in different hardware platforms that range from server-grade platforms to edge devices such as smartphones . Optimal network architectures need to be designed to meet the requirement for a target deployment scenario . However , naively designing a specialized architecture for each scenario is very expensive as it requires to fully retrain the model each time . This is an excessively expensive process in terms of the required machine learning expertise , time , energy and CO2 emission . Recently , researchers have proposed efficient methods which are based on training a super network only once . Then , for a specific deployment scenario , a sub-network is sampled from the super-network that meet the deployment constraints with the best accuracy . The weight of the sampled network is shared with original super network , hence retraining is not required . Once-for-all ( OFA ) ( Cai et al. , 2020 ) is among the first methods proposed to tackle this problem . The OFA method trains a once-for-all network that jointly optimizes the accuracy of a large number of subnetworks ( more than 1019 ) sampled from the once-for-all network . Each sub-network is selected from the once-for-all network where layer depths , channel widths , kernel sizes and input resolution are scaled independently . Such scaling provides a family of CNNs networks with different computation and representation power to flexibly support deployment under diverse platforms and configurations . With this massive search space , OFA co-trains all the sub-networks by a complex four-stage progressive training process which is prohibitively expensive and costs around 1200 GPU hours . Compound OFA ( CompOFA ) ( Sahni et al. , 2020 ) builds upon the original OFA by shrinking the design space of possible sub-networks . This is done by only considering networks whose dimensions are coupled . This reduces the number of possible models by 17 orders of magnitudes , from 1019 down to 243 . Sahni et al . ( 2020 ) demonstrate that this smaller design space is sufficient , as most warmup phase in-place distillation teacher training compound training teacher training elastic kernel elastic depth elastic widthOFA compOFA fOFA Figure 1 : Comparison of training schedules for OFA , CompOFA , and fOFA . Length on the horizontal axis is proportional to the number of epochs in each phase . For OFA , “ elastic kernel , ” “ elastic width , ” and “ elastic depth ” are the phases of training specified in Cai et al . ( 2020 ) that are not used in CompOFA and fOFA . sub-networks in the original OFA design space are far from the optimal accuracy-latency frontier . With this smaller space , the training procedure can be simplified as well as these suboptimal subnetworks are no longer influencing the training process . CompOFA reduces the four stages of the original OFA process to two stages , and this optimization speeds up the training time of CompOFA by a factor of 2× over OFA . However , 2× faster than OFA ’ s 1200 GPU-hours is still 600 GPU-hours . Even with this significant improvement , the training cost remains vary expensive , especially when effects on the environment are considered ( Strubell et al. , 2019 ) . While some of this cost can be mitigated by improvements in hardware efficiency and the continued development of specialized platforms for training CNNs , algorithmic enhancements still have a large role to play . While CompOFA greatly simplifies the progressive shrinking training procedure used in OFA , it is still dependent on pre-training a supernetwork to act as a teacher for the sub-network co-training process , which uses knowledge distillation Hinton et al . ( 2015 ) . Due to the optimizations in the co-training process , training the supernetwork in CompOFA requires more than half ( 180 out of 330 ) of the total training epochs . In this work , we propose several optimizations to the once-for-all training process that produces a one-stage training algorithm for fast and efficient neural architecture search . The key features of our method are : • We co-train all of the sub-networks from scratch without pre-training a teacher network , using the concept of in-place distillation ( Yu et al. , 2020 ) . The largest network we train using in-place distillation is smaller than the pre-trained teacher network used in Cai et al . ( 2020 ) and Sahni et al . ( 2020 ) . • During the co-training process , we develop an upper-attentive sampling method which always sample the full-sized sub-net at each iteration to help co-train the rest sub-networks . • Before co-training , we use an upper-attentive warmup technique which trains only the fullsized sub-net for a few epochs before co-training to further improve the performance . • With these optimizations , we can decrease the number of sampled sub-networks in each iteration of training , further improving performance . The benefits of our proposed fast OFA ( fOFA ) method are shown in Figure 1 . Furthermore , since our method has only a single stage , we can easily increase the training time and improve on the accuracy of previous methods while still requiring less training time . The rest of this paper is organized as follows . We describe related work in more detail in Section 2 and illustrate our method in depth in Section 3 . We report on our experimental results in Section 4 and finish the paper with conclusions in 6 . 2 RELATED WORK . Neural architecture search ( NAS ) aims to automatically find the optimal network architecture given hardware constraints , such as FLOPs or latency . Early NAS works mainly adapted reinforcement learning ( Zoph et al. , 2018 ; Zoph & Le , 2016 ) , evolutionary search ( Real et al. , 2019 ; 2017 ) , or sparse connection learning ( Kim et al. , 2018 ) to sample different architectures . However , each sampled architecture needed to be trained from scratch , resulting in a huge and intractable computing cost . More recent NAS works greatly reduce the cost by training an over-parameterized network named a super-network , and then sample various sub-networks which share the weights with the super-network . Such super-network-based methods can be further divided into two main categories as follows : 2.1 TWO-STAGE TRAINING . The main idea of the two-stage training methods ( Berman et al. , 2020 ; Bender et al. , 2018 ; Brock et al. , 2017 ; Guo et al. , 2019 ; Liu et al. , 2018 ; Pham et al. , 2018 ) is that after searching for the best architectures in the first stage of training , the best architectures then have to be retrained from scratch to obtain a final model . Generally , a single two-stage search experiment can only target a single resource budget or a narrow range of resource budgets at a time , which is inefficient . 2.2 ONE-STAGE TRAINING . To alleviate the inefficiency of two-stage training , Once-For-All ( OFA ) ( Cai et al. , 2020 ) was proposed to jointly train various sub-networks of the super-network in a single stage . By doing so , the sub-networks could be directly deployed into different hardware platform without retraining . However , to support an extremely large number of sub-networks ( i.e. , 1019 ) , such one-stage training involves multi-steps to gradually add more sub-networks by using the proposed progressive shrinking technique . Moreover , OFA also needs to train a single full-sized teacher network ( same size as super-network ) from scratch firstly to guide the training of the sub-networks by using knowledge distillation . Thus , due to the complex training procedure , OFA still suffers from a prohibitive training cost , requiring around 1200 GPU hours . More recently , inspired by studies on neural network design spaces ( Tan & Le , 2019 ; Radosavovic et al. , 2020 ) , CompOFA ( Sahni et al. , 2020 ) proposes a compound sub-network scaling method , which couples the depth and width configuration of the sampled sub-networks to constrain the search space , reducing the space to only 243 number of sub-networks without losing accuracy . Although CompOFA achieves a 2× training cost reduction compared with OFA , it follows the similar training procedure that also needs to train a teacher model from scratch . A similar approach is investigated in Yang et al . ( 2020 ) , which couples network width and input resolution into a single mutual learning framework . In addition , bigNAS ( Yu et al. , 2020 ) proposed replacing OFA ’ s multi-step training by a single step , namely one-shot NAS , challenging the usual practice of progressive training in OFA . The idea of the one-shot NAS is to jointly train sub-networks from scratch directly by using the sandwich rule and in-place distillation techniques proposed for slimmable networks ( Yu et al. , 2018 ) . Based on bigNAS , Wang et al . ( 2020 ) points out the unnecessary updates on sub-optimal models in one-stage training , and uses attention mechanisms to push the Pareto front . However , the primary objective of these two works is to obtain better accuracy , thus still suffer from high training cost , e.g. , bigNAS needs over 2300 TPU hours for O ( 1012 ) models . In addition , ( Li et al. , 2021 ) works to improve the trade-off between accuracy and computation complexity based for slimmable networks by introducing a dynamic gating mechanism and in-place ensemble bootstrapping to increase training stability .. However , it requires a one more gating training step , resulting in larger training cost . Our approach also follows the direction of the one-shot model . Differentiating from bigNAS , the primary objective of this work is to reduce the training cost without loss of Pareto optimality under the design space of OFA and CompOFA . 3 METHODS . 3.1 BUILDING THE SEARCH SPACE . A neural network N is a function that takes an input set X and generates a set of outputs δ ( N , X ) . In this work , we focus on a fixed input set ( i.e. , ImageNet ) , and thus write the network output as δ ( N ) . In the supervised learning setting , the performance of the neural network is evaluated against a set of labels YD . Following the standard practice in neural architecture search , we limit our neural network space to the set of architectures that consists of a sequence of blocks B1 , B2 , . . . , Bm , where m = 5 is a typical value . Each block is based on the inverted residual in the architecture space of MobileNetV3 ( Howard et al. , 2019 ) . A block is parameterized by three dimensions : the depth ( number of layers in the block ) D , the width ( channel expansion ratio ) W , and the convolution kernel size K. This search space is illustrated in Figure 2 . To reduce the size of the search space , we use the same coupling heuristic as CompOFA ( Sahni et al. , 2020 ) ; that is , if there are n choices for the depth dimension and n choices for the width dimension , we sample the ith largest depth wi whenever we sample the ith largest depth di for each layer in the block . While OFA uses an elastic kernel that allows for different kernel sizes within blocks , we follow CompOFA and use a fixed kernel size within each block . We call the network where the values of K , D , and W are each their largest possible value the full-sized network or super-net , and the network created by any other choice of these values a sub-network . As in CompOFA , we choose three possible values for D ∈ { 2 , 3 , 4 } and three possible values for W ∈ { 3 , 4 , 6 } and fix the kernel size to that of ( Howard et al. , 2019 ) , that is , K = 3 in the first , third , and fourth blocks , and K = 5 in the second and fifth blocks . Thus , with five blocks , we have 35 = 243 models in our search space . In neural architecture search , the input resolution can vary as well , up to a maximum size of 224 × 224 for ImageNet ( Deng et al. , 2009 ) . In this work , we use an elastic resolution , where input images are resized to be square with dimension in the set { 128 , 160 , 192 , 224 } .

Summary: In this manuscript, the authors propose a new framework for training the once-for-all (OFA) networks. This framework uses an in-place distillation schema to replace the teacher training in the former OFA training fashion and makes the overall training time faster. The paper is overall well-written with good performance.