AtomNAS: Fine-Grained End-to-End Neural Architecture Search

1 INTRODUCTION . Human-designed neural networks are already surpassed by machine-designed ones . Neural Architecture Search ( NAS ) has become the mainstream approach to discover efficient and powerful network structures ( Zoph & Le ( 2017 ) ; Pham et al . ( 2018 ) ; Tan et al . ( 2019 ) ; Liu et al . ( 2019a ) ) . Although the tedious searching process is conducted by machines , humans still involve extensively in the design of the NAS algorithms . Designing of search spaces is critical for NAS algorithms and different choices have been explored . Cai et al . ( 2019 ) and Wu et al . ( 2019 ) utilize supernets with multiple choices in each layer to accommodate a sampled network on the GPU . Chen et al . ( 2019b ) progressively grow the depth of the supernet and remove unnecessary blocks during the search . Tan & Le ( 2019a ) propose to search the scaling factor of image resolution , channel multiplier and layer numbers in scenarios with different computation budgets . Stamoulis et al . ( 2019a ) propose to use different kernel sizes in each layer of the supernet and reuse the weights of larger kernels for small kernels . Howard et al . ( 2019 ) ; Tan & Le ( 2019b ) adopts Inverted Residuals with Linear Bottlenecks ( MobileNetV2 block ) ( Sandler et al. , 2018 ) , a building block with light-weighted depth-wise convolutions for highly efficient networks in mobile scenarios . However , the proposed search spaces generally have only a small set of choices for each block . DARTS and related methods ( Liu et al. , 2019a ; Chen et al. , 2019b ; Liang et al. , 2019 ) use around 10 different operations between two network nodes . Howard et al . ( 2019 ) ; Cai et al . ( 2019 ) ; Wu et al . ( 2019 ) ; Stamoulis et al . ( 2019a ) search the expansion ratios in the MobileNetV2 block but still limit them to a few discrete values . We argue that search space of finer granularity is critical to find ∗ This work was done during the internship program at Bytedance . optimal neural architectures . Specifically , the searched building block in a supernet should be as small as possible to generate the most diversified model structures . We revisit the architectures of state-of-the-art networks ( Howard et al . ( 2019 ) ; Tan & Le ( 2019b ) ; He et al . ( 2016 ) ) and discover a commonly used building structure : convolution - channel-wise operation - convolution . We reinterpret this building structure as an ensemble of computationally independent blocks , which we call atomic blocks . As the minimum search unit , the atomic block constitutes a much larger and more fine-grained search space , within which we are able to search for mixed operations ( e.g. , convolutions with different kernel sizes and their channel numbers ) . For the efficient exploration of the new search space , we propose a NAS framework named AtomNAS which applies network pruning techniques to architecture search . Specifically , we start from an initial large supernet and rewrite every convolution - channel-wise operation - convolution structure of it in the form the weighted sum of atomic blocks ; the weights reflect the contribution of the atomic blocks to the network capacity and are called importance factors . For each atomic block , a penalty term in proportion to its FLOPs is enforced on its importance factor ; effectively , the penalty makes AtomNAS favor atomic blocks with less FLOPs . By minimizing the combination of the original network loss and the total penalty on the weights , AtomNAS is able to learn both the parameters of the network and the weights of the atomic blocks . At the end of the learning , atomic blocks with very small weights ( e.g. , < 0.001 ) are removed from the network and we obtain the final network which has fewer FLOPs . Since the pruned atomic blocks have little contribution to the network output due to their negligible weights , the final network does not need to be retrained or finetuned . Training on the large supernet is computationally demanding . We observe that for many pruned atomic blocks , their weights diminish at the early stage of learning and never “ revive ” throughout the rest of learning . We propose a dynamic network shrinkage technique which removes those atomic blocks on the fly and greatly reduces the run time of AtomNAS . In our experiment , our method achieves 75.9 % top-1 accuracy on ImageNet dataset around 360M FLOPs , which is 0.9 % higher than state-of-the-art model ( Stamoulis et al. , 2019a ) . By further incorporating additional modules , our method achieves 77.6 % top-1 accuracy . It outperforms MixNet by 0.6 % using 363M FLOPs , which is a new state-of-the-art under the mobile scenario . In summary , the major contributions of our work are : 1 . We design a fine-grained search space which includes the exact number of channels and mixed operations ( e.g. , combination of different convolution kernels ) . 2 . We propose an NAS framework , AtomNAS . Within the framework , an efficient end-to-end NAS algorithm is proposed which can simultaneously search the network architecture and train the final model . No finetuning is needed after the algorithm finishes . 3 . With the proposed search space and AtomNAS , we achieve state-of-the-art performance on ImageNet dataset under mobile setting . 2 RELATED WORK . 2.1 NEURAL ARCHITECTURE SEARCH . Recently , there is a growing interest in automated neural architecture design . Reinforce learning based NAS methods ( Zoph & Le , 2017 ; Tan et al. , 2019 ; Tan & Le , 2019b ; a ) are usually computational intensive , thus hampering its usage with limited computational budget . To accelerate the search procedure , ENAS ( Pham et al. , 2018 ) represents the search space using a directed acyclic graph and aims to search the optimal subgraph within the large supergraph . A training strategy of parameter sharing among subgraphs is proposed to significantly increase the searching efficiency . The similar idea of optimizing optimal subgraphs within a supergraph is also adopted by Liu et al . ( 2019a ) ; Jin et al . ( 2019 ) ; Xu et al . ( 2020 ) ; Wu et al . ( 2019 ) ; Guo et al . ( 2019 ) ; Cai et al . ( 2019 ) . Stamoulis et al . ( 2019a ) ; Yu et al . ( 2020 ) further share the parameters of different paths within a block using super-kernel representation . A prominent disadvantage of the above methods is that their coarse search spaces only support selecting one out of a set of choices ( e.g. , selecting one kernel size from { 3 , 5 , 7 } ) . MixNet tries to benefit from mixed operations by using a predefined set of mixed operations { { 3 } , { 3 , 5 } , { 3 , 5 , 7 } , { 3 , 5 , 7 , 9 } } , where the channels are equally distributed among different kernel sizes . Due to this limitation , it is difficult to learn optimal architectures under computational resource constraints . On the contrary , our method takes advantage of the fine-grained search space and is able to search for more flexible network architectures satisfying various resource constraints . The fine-grained search space proposed in this paper is exponentially larger than previous search space . For reference , the total number of possible structures within the experiment is around 10162 , compared with 1021 for FBNet . Recently , to improve the final performance of the searched architectures , Yu et al . ( 2020 ) utilizes knowledge distillation which is orthogonal to our method . It could be easily integrated into our method by Eq . ( 5 ) thanks to the end-to-end learning paradigm of our method . 2.2 NETWORK PRUNING . Assuming that many parameters in the network are unnecessary , network pruning methods start from a computation-intensive model , identify the unimportant connections and remove them to get a compact and efficient network . Early method ( Han et al. , 2016 ) simultaneously learns the important connections and weights . However , non-regularly removing connections in these works makes it hard to achieve theoretical speedup ratio on realistic hardwares due to extra overhead in caching and indexing . To tackle this problem , structured network pruning methods ( He et al. , 2017b ; Liu et al. , 2017 ; Luo et al. , 2017 ; Ye et al. , 2018 ; Gordon et al. , 2018 ) are proposed to prune structured components in networks , e.g . the entire channel and kernel . In this way , empirical acceleration can be achieved on modern computing devices . Liu et al . ( 2017 ) ; Ye et al . ( 2018 ) ; Gordon et al . ( 2018 ) encourage channel-level sparsity by imposing the L-1 regularizer on the channel dimension , which is also used by our method . Recently , Liu et al . ( 2019b ) show that in structured network pruning , the learned weights are unimportant . This suggests structured network pruning is actually a neural architecture search focusing on channel numbers . Our method jointly searches the channel numbers and a mix of operations , which is a much larger search space . 3 ATOMNAS . We formulate our neural architecture search method in a fine-grained search space with the atomic block used as the basic search unit . An atomic block is comprised of two convolutions connected by a channel-wise operation . By stacking atomic blocks , we obtain larger building blocks ( e.g . residual block and MobileNetV2 block proposed in a variety of state-of-the-art models including ResNet , MobileNet V2/V3 ( He et al. , 2016 ; Howard et al. , 2019 ; Sandler et al. , 2018 ) . In Section 3.1 , We first show larger network building blocks ( e.g . MobileNetV2 block ) can be represented by an ensembles of atomic blocks . Based on this view , we propose a fine-grained search space using atomic blocks . In Section 3.2 , we propose a resource-aware atomic block selection method for end-to-end architecture search . Finally , we propose a dynamic network shrinkage technique in Section 3.3 , which greatly reduces the search cost . 3.1 FINE-GRAINED SEARCH SPACE . Under the typical block-wise NAS paradigm ( Tan et al. , 2019 ; Tan & Le , 2019b ) , the search space of each block in a neural network is represented as the Cartesian product C = ∏ i=1 Pi , where each Pi is the set of all choices of the i-th configuration such as kernel size , number of channels and type of operation . For example , C = { conv , depth-wise conv , dilated conv } × { 3 , 5 } × { 24 , 32 , 64 , 128 } represents a search space of three types of convolutions by two kernel sizes and four options of channel number . A block in the resulting model can only pick one convolution type from the three and one output channel number from the four values . This paradigm greatly limits the search space due to the few choices of each configuration . Here we present a more fine-grained search space by decomposing the network into smaller and more basic building blocks . We denote f c ′ , c ( X ) as a convolution operator , whereX is the input tensor and c , c′ are the input and output channel numbers respectively . A wide range of manually-designed and NAS architectures share a structure that joins two convolutions by a channel-wise operation : Y = ( f c ′′ , c′ 1 ◦ g ◦ f c′ , c 0 ) ( X ) ( 1 ) where g is a channel-wise operator . For example , in VGG ( Simonyan & Zisserman , 2015 ) and a Residual Block ( He et al. , 2016 ) , f0 and f1 are convolutions and g is one of Maxpool , ReLU and BN-ReLU ; in a MobileNetV2 block ( Sandler et al. , 2018 ) , f0 and f1 are point-wise convolutions and g is depth-wise convolution with BN-ReLU in the MobileNetV2 block . Eq . ( 1 ) can be reformulated as follows : Y = c′∑ i=1 ( f c ′′,1 1 [ i , : ] ◦ g [ i , : ] ◦ f 1 , c 0 [ : , i ] ) ( X ) , ( 2 ) where f1 , c0 [ : , i ] is the i-th convolution kernel of f0 , g [ i , : ] is the operator of the i-th channel of g , and { f c ′′,1 1 [ i , : ] } c ′ i=1 are obtained by splitting the kernel tensor of f1 along the the input channel dimension . Each term in the summation can be seen as a computationally independent block , which is called atomic block . Fig . ( 1 ) demonstrate this reformulation . By determining whether to keep each atomic block in the final model individually , the search of channel number c′ is enabled through channel selection , which greatly enlarges the search space . This formulation also naturally includes the selection of operators . To gain a better understanding , we first generalize Eq . ( 2 ) as : Y = c′∑ i=1 ( f c ′′,1 1i ◦ gi ◦ f 1 , c 0i ) ( X ) . ( 3 ) Note the array indices i are moved to subscripts . In this formulation , we can use different types of operators for f0i , f1i and gi ; in other words , f0 , f1 and g can each be a combination of different operators and each atomic block can use different operators such as convolutions with different kernel sizes . Formally , the search space is formulated as a supernet which is built based on the structure in Eq . ( 1 ) ; such structure satisfies Eq . ( 3 ) and thus can be represented by atomic blocks ; each of f0 , f1 and g is a combination of operators . The new search space includes some state-of-the-art network architectures . For example , by allowing g to be a combination of convolutions with different kernel sizes , the MixConv block in MixNet ( Tan & Le , 2019b ) becomes a special case in our search space . In addition , our search space facilitates discarding any number of channels in g , resulting in a more fine-grained channel configuration . In comparison , the channel numbers are determined heuristically in Tan & Le ( 2019b ) .

Summary: This basic idea of this paper is to decompose the common building blocks of large network into atomic blocks, which equips NAS with more fine-grained search space. What's more, the authors propose a resource-aware search to reduce the computation and dynamically shrinkage the model to accelerate the learning. Retraining the final network is no longer needed. They achieve state of art on ImageNet under several complexity constraints.