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: [Summary] This paper proposes a channel-wise neural architecture search (NAS) approach. The NAS search algorithm is similar to previous one-shot NAS, but the search space is channel-wise: each channel has it’s own kernel size, which is quite novel and interesting. Results are strong in terms of FLOPS and parameters.