Efficient Differentiable Neural Architecture Search with Model Parallelism

1 INTRODUCTION . Neural architecture search ( NAS ) has revolutionized architecture designs of deep learning from manually to automatically in various applications , such as image classification ( Zoph & Le , 2016 ) and semantic segmentation ( Liu et al. , 2019a ) . Reinforcement learning ( Zoph & Le , 2016 ; Zoph et al. , 2018 ; Pham et al. , 2018 ) , evolutionary algorithms ( Real et al. , 2017 ; 2019 ) , and differentiable algorithms ( Liu et al. , 2019b ; Cai et al. , 2019 ) have been applied to discover the optimal architecture from a large search space of candidate network structures . Supernets ( Zoph et al. , 2018 ; Pham et al. , 2018 ) comprising all possible networks reduce search spaces from complete network architectures to cell structures . Recent acceleration techniques of differentiable NAS ( Xie et al. , 2019 ; Yao et al. , 2020 ; Chen et al. , 2019 ; Xu et al. , 2020 ) further diminish search costs to affordable computation overheads ( e.g. , half GPU day ) . Prior work ( Xu et al. , 2020 ) randomly samples partial channels of intermediate feature maps in the mixed operations . However , supernets of differentiable NAS consume gigantic GPU memory , which constrains NAS from using large batch sizes and imposes restrictions on supernet architectures ’ complexity . For example , NAS determines networks in shallow supernets ( e.g. , 8 layers ) for deep compact networks ( e.g. , 20 layers ) . The cell structures are also required to remain identical for the same type of cells . Data parallelism can increase the search efficiency of NAS by using large batch sizes , such as SNAS ( Xie et al. , 2019 ) , but it requires supernet complexity low enough to fit in a single GPU . In contrast , model parallelism can parallelize complex supernets , which distributes partial models to multiple devices . Nevertheless , model parallelism suffers from low hardware utilization . Only one device executes its model partition , while other devices stay idle . How to take advantage of multiple GPUs for large supernets efficiently is an open problem . 1Search and evaluation code are released at link In this paper , we propose a simple and efficient solution , binary neural architecture search ( NASB ) using consecutive model parallel ( CMP ) , to tackle the above limitations . Specifically , supernets have two forward and two backward phases to learn architecture parameters and network weights . CMP distributes several sub-tasks split from the four phases in multiple GPUs and executes the sub-tasks of all forward/backward phases together . Figure 1 illustrates that sub-tasks of forward/backward phases will be overlapped to reduce waiting cycles . Nevertheless , CMP consumes large GPU memory due to two computation graphs existing at the same time . Thus , we introduce NASB to declines GPU memory occupation . NASB utilizes binary and sparse architecture parameters ( 1 or 0 ) for mixed operations . It excludes inactive operations in the computation graph and computes feature maps of inactive operations for architecture gradients during the back-propagation . In this way , NASB-CMP can increase hardware utilization of model parallelism with efficient GPU memory in differentiable NAS . In our experiments on CIFAR-10 , NASB-CMP runs 1.2× faster than using model parallel and pipeline parallel , TorchGPipe ( Kim et al. , 2020 ) in a server with 4 GPUs 2 . It can achieve the test error of 2.53 ± 0.06 % by searching for only 1.48 hours . Our contribution can be summarized as follows : • NASB-CMP is the first NAS algorithm that can parallelize large supernets with large batch sizes . We analyze the acceleration ratio between CMP and traditional model parallelism . Even though complex supernets ( e.g. , large layers and different cell structures ) will not boost NAS performance , NASB-CMP paves the way to explore the supernet architecture design in the future . • NASB utilizes binary architecture parameters and extra architecture gradients computation to reduce GPU usage . It can save memory consumption by accepting twice batch sizes larger than the other memory saving algorithm , PC-DARTS ( Xu et al. , 2020 ) . • We fairly compare NASB-CMP with state-of-the-art differentiable NAS in the same hardware and search space . Extensive experiments show that NASB-CMP can achieve competitive test error in short search time . 2NVIDIA GTX 1080 Ti . 2 METHODOLOGY . We first describe the fundamental concepts of one-shot neural architecture search ( NAS ) in Section 2.1 . We then portray the consecutive model parallel to enhance NAS search efficiency in multiple devices in Section 2.2 . Finally , we explain how we binarize the architectural weights and compute their gradients to cut down the GPU memory consumption in Section 2.3 . 2.1 ONE-SHOT NEURAL ARCHITECTURE SEARCH . One-shot neural NAS ( Zoph et al. , 2018 ) is built on a supernet ( a.k.a . meta graph ) in which we stack normal cells and reduce cells sequentially in Figure 2 ( a ) . Normal cells are analogous to convolutional layers to extract images features . Reduce cells are equivalent to pooling layers to reduce the spatial dimension of feature maps . All normal cells share the same structure , but each cell still has its network weights . So do all reduce cells . One-shot approaches are required to design two cell structures instead of complete neural networks . Figure 2 ( b ) illustrates one popular cell structure ( Pham et al. , 2018 ) , an N -node directed-acyclic-graph ( DAG ) with total edges E , not counting the “ concat ” node . In the h-th cell , the first two nodes are the ( h − 2 ) -th and ( h − 1 ) -th cells having no inbound edges . The other nodes accept previous nodes whose index is lower than the current index . Total edges E ( red lines of Figure 2 ( b ) ) is ( N +1 ) ( N −2 ) /2 . We denote the h-th cell ’ s output as yh = concat ( nj ) , where 2 ≤ j ≤ N − 1 and nj is a DAG node signified in Eq . 1. nj = yh−2 , if j = 0 , yh−1 , j = 1 , ∑ i < j mO ( ni ) , 2 ≤ j ≤ N − 1 . ( 1 ) A mixed operation mO is the edge between node i and j in the DAG . Let O be a set of candidate operations ( e.g. , convolution , pooling , identity , zero ) and A ∈ RE×|O| be a matrix of architecture parameters . Eq . 2 formulates the mixed operation mO from node i to j as the weighted sum of all operations ok ( Liu et al. , 2019b ) . mjO ( ni ) = |O|∑ k=1 Ae , kok ( ni ) , j ≥ 2 , i < j , ( 2 ) where e = ( j+1 ) ( j− 2 ) /2+ i is the edge index . The mixed operations transform the cell structure search to the problem of learning two matrices , AN and AR , for the normal and reduce cell . Given that Lval and Ltrain is the loss function L beyond a training and validation dataset , respectively . Let A comprise AN and AR . Mathematically , one-shot NAS can be formulated in the following optimization problem , minA Lval ( w ∗ , A ) s.t . w∗ = argmin w Ltrain ( w , A ) . ( 3 ) NAS leverages the validation performance to choose well-trained networks that outperform others . After training A , we derive the compact network by pruning unused operations in the supernet . Since the whole paper follows the image classification setting ( Liu et al. , 2019b ; Cai et al. , 2019 ) , we assume each node is assigned two inputs and two operations . And we prune node inputs of cells of the supernet by the largest two values of A associated with that node . For simplicity , we use A in replace of A in the following discussion . 2.2 CONSECUTIVE MODEL PARALLEL . Data parallelism can scale up supernets with large batch sizes , but it can not handle large supernets ( e.g. , deep supernets with different cell structures ) . Model parallelism ( MP ) is able to amortize such large supernets across multiple GPUs , but its hardware utilization is low . MP would generate unwanted waiting cycles across devices . Figure 1 displays that the first device becomes idle until the second device finishes its forward and backward phases . The parallelization gets worse as we use large available GPUs . Motivated by pipeline parallelism ( Huang et al. , 2019 ) , we propose consecutive model parallel ( CMP ) to decrease GPU idle time . Let FA and BA signify the forward and backward phase to update A , and Fw and Bw be two phases to update w. CMP divides the four phases into several sub-tasks and performs sub-tasks of FA and Fw consecutively , followed by sub-tasks of Bw and BA . Figure 1 illustrates that the execution order change by CMP overlaps sub-tasks without waiting for others to finish . Given the number of available GPUs M , Eq . 4 reveals the ratio of execution time between CMP and MP in theory . Time of CMP Time of MP = 1 M [ 4M − 2 ( M − 1 ) ] 4 = 1− M − 1 2M . ( 4 ) We assume FA , BA , Fw , and Bw take the same time unit . MP will complete an iteration in 4 units . For CMP , the total sub-tasks is 4M , and 2 ( M − 1 ) sub-tasks can be overlapped . If a sub-task takes 1/M ideally , CMP will finish an iteration in 1/M ( 4M − 2 ( M − 1 ) ) units . According to Eq . 4 , CMP with two devices could reduce ( 2-1 ) / ( 2 * 2 ) =25 % time from MP . In practice , Experiment 3.1 demonstrates that NASB-CMP runs 1.2× faster than model parallelism without sacrificing test error . The theoretical value for 4 GPU is 1.6 ( or reduce 37.5 % time ) . We believe communication overhead and uneven model balance cause the deviation . Communication overhead comes from the intermediate tensors transfer from one to another GPU when models are split into different GPUs . Moreover , the main thread is responsible for loading data and backward propagation . The GPU with the main thread always consumes the most GPU memory , which causes uneven model balance . CMP is a general model parallel approach for any existing differentiable NAS algorithm . However , runningBA andBw consecutively asks for two computation graphs , which doubles GPU utilization and deteriorates CMP efficiency . To address the problem of great GPU consumption , we introduce a memory-efficient NAS to CMP , called binary neural architecture search ( NASB ) . 2.3 BINARY NEURAL ARCHITECTURE SEARCH . Binary neural architecture search ( NASB ) harnesses binary mixed operations mBO ( Yao et al. , 2020 ) that convert the real-valued A into sparse binary matrix G , as illustrated in Figure 2 . Among rows Ae , : associate node j , mBO enforces the two largest elements to 1 ( active ) and the rest elements to 0 ( inactive ) . The row indexes of active elements indicate selected edges to node j , while column indexes indicate chosen operations . Notice that NASB does not directly multiply G with candidate operations in Eq . 5 . Instead , NASB constructs a set of active operations O ( active ) based on active elements in G. Only those active operations oa ∈ O ( active ) are included in the forward phase . This technique could stop inactive operations being stored in the computation graph and decrease roughly Algorithm 1 : NASB - Consecutive Model Parallel 1 : Initialize architecture weights A and network weights w 2 : while not stopped do 3 : Gt = binarize ( At ) 4 : Create mBO using Gt and Eq . 5 5 : Compute Lvalid ( wt , Gt ) and Ltrain ( wt , Gt ) consecutively // model parallel 6 : Compute∇wLtrain ( wt , Gt ) and∇ALvalid ( wt , Gt ) consecutively // model parallel 7 : Update wt+1 by descending∇wLtrain ( wt , Gt ) 8 : Update At+1 by descending∇ALvalid ( wt , Gt ) 9 : end while |O| times GPU memory compared to using the multiplication by G. mBO ( ni ) = |O|∑ k=1 Ge , kok ( ni ) = oa ( ni ) . ( 5 ) NASB computes gradients of network weights w using standard back-propagation in the supernet . For the gradients of A , NASB estimates ∂L/∂A approximately by ∂L/∂G : ∂L ∂Ae , k = ∂L ∂mO ∂mO ∂Ae , k ≈ ∂L ∂mBO ∂mBO ∂Ge , k = ∂L ∂mBO × ok ( n ) = ∂L ∂Ge , k . ( 6 ) Eq . 6 states that gradients of elements in A come from ∂L/∂mBO × ok ( n ) . However , inactive operations are not in the computation graph . NASB saves inputs of inactive operations n in PyTorch Context that is used for backward computation . During the backward phase , NASB will compute inactive operations ok′ ( n ) on the fly and multiply the results with the ∂L/∂mBO . Apart from saving unneeded GPU FLOPS and memory , mBO can avoid performance bias between supernets and compact networks . Supernets using mO assume that the performance of supernets can represent derived compact networks , but non-linear operations ( e.g. , ReLU-Conv-BN ) break the representation that causes performance bias ( Xie et al. , 2019 ) . Instead , the sparse matrix of mBO activates one operation . The performance of supernets during the search is only for one compact network . Thus , NASB can mitigate the bias caused by non-linear operations . Algorithm 1 describes how CMP works with NASB . Note that NASB-CMP does not update any parameter ( including A and w ) until FA , BA , Fw , and Bw complete . Ltrain will use the current binary architecture matrix Gt rather than updated Gt+1 , which is the major difference from the alternate algorithm ( See Appendix A ) . Experiment 3.2 demonstrates NASB could save substantial GPU memory than PC-DARTS ( Xu et al. , 2020 ) , which reduces GPU memory by partial channels of feature maps in mixed operations . Comparison with other methods . NASP ( Yao et al. , 2020 ) binarizes A based on A itself , while ProxylessNAS ( Cai et al. , 2019 ) binarizes A based on the softmax results of A . The two binarization approaches are equivalent , but how they handle binary mixed operations ( Eq . 5 ) is different . NASP multiplies G with all operations ( i.e. , saving active and inactive operations in the computation graph ) . ProxylessNAS selects two sampled operations ( paths ) in the computation graph according to multinomial distribution . NASB utilizes the same binarization as NASP but only keeps one active operation in the computation graph according to G .

Summary: This paper provides the interesting method that leverages GPU memory resources more efficiently for supernet (meta-graph) of differentiable NAS. For this, this paper proposes binary neural architecture search and consecutive model parallel (CMP). CMP parallelizes one supernet with multiple GPUs, which allows NAS model to use larger batch size and search space. Additionally, this paper improves neural architecture search speed and hardware utilization with waiting cycles reduction by dividing forward/backward phases into several sub-tasks and executing the same type of sub-tasks. The proposed method shows 1.2x faster search time compared with other model parallel methods and the highest performance among differentiable NAS methods in the experiment section.