Densely Connected Search Space for More Flexible Neural Architecture Search

Introduction

In recent years, neural architecture search (NAS) has demonstrated great successes in designing neural architectures automatically and achieved remarkable performance gains in various tasks such as image classification , semantic segmentation and object detection . NAS has been a critically important topic for architecture designing.

In NAS research, the search space plays a crucial role that constrains the architectures in a prior-based set. The performance of architectures produced by NAS methods is strongly associated with the search space definition. A more flexible search space has the potential to bring in architectures with more novel structures and promoted performance. We revisit and analyze the search space design in most previous works . For a clear illustration, we review the following definitions. Block denotes a set of layers/operations in the network which output feature maps with the same spatial resolution and the same width (number of channels). Stage denotes a set of sequential blocks whose outputs are under the same spatial resolution settings. Different blocks in the same stage are allowed to have various widths. Many recent works stack the inverted residual convolution modules (MBConv) defined in MobileNetV2 to construct the search space. They search for different kernel sizes and expansion ratios in each MBConv. The depth is searched in terms of layer numbers in each block. The searched networks with MBConvs show high performance with low latency or few FLOPs.

In this paper, we aim to perform NAS in a more flexible search space. Our motivation and core idea are illustrated in Fig. 1. As the upper part of Fig. 1 shows, the number of blocks in each stage and the width of each block are set manually and fixed during the search process. It means that the depth search is constrained within the block and the width search cannot be performed. It is worth noting that the scale (depth and width) setting is closely related to the performance of a network, which has been demonstrated in many previous theoretical studies and empirical results . Inappropriate width or depth choices usually cause drastic accuracy degradation, significant computation cost, or unsatisfactory model latency. Moreover, we find that recent works manually tune width settings to obtain better performance, which indicates the design of network width demands much prior-based knowledge and trial-and-error.

We propose a densely connected search space to tackle the above obstacles and name our method as DenseNAS. We show our novelly designed search space schematically in the bottom part of Fig. 1. Different from the search space design principles in the previous works , we allow more blocks with various widths in one stage. Specifically, we design the routing blocks to construct the densely connected super network which is the representation of the search space. From the beginning to the end of the search space, the width of the routing block increases gradually to cover more width options. Every routing block is connected to several subsequent ones. This formulation brings in various paths in the search space and we search for the best path to derive the final architecture. As a consequence, the block widths and counts in each stage are allocated automatically. Our method extends the depth search into a more flexible space. Not only the number of layers within one block but also the number of blocks within one stage can be searched. The block width search is enabled as well. Moreover, the positions to conduct spatial down-sampling operations are determined along with the block counts search.

We integrate our search space into the differentiable NAS framework by relaxing the search space. We assign a probability parameter to each output path of the routing block. During the search process, the distribution of probabilities is optimized. The final block connection paths in the super network are derived based on the probability distribution. To optimize the cost (FLOPs/latency) of the network, we design a chained estimation algorithm targeted at approximating the cost of the model during the search.

Our contributions can be summarized as follows.

We propose a densely connected search space that enables network/block widths search and block counts search. It provides more room for searching better networks and further reduces expert designing efforts.

We propose a chained cost estimation algorithm to precisely approximate the computation cost of the model during search, which makes the DenseNAS networks achieve high performance with low computation cost.

In experiments, we demonstrate the effectiveness of our method by achieving SOTA performance on the MobileNetV2 -based search space. Our searched network achieves 75.3% accuracy on ImageNet with only 361MB FLOPs and 17.9ms latency on a single TITAN-XP.

DenseNAS can further promote the ImageNet classification accuracies of ResNet-18, -34 and -50-B by 1.5%, 0.5% and 0.3% with 200M, 600M, 680M FLOPs and 1.5ms, 2.4ms, 6.1ms latency reduction respectively.

Related Work

NASNet is the first work to propose a cell-based search space, where the cell is represented as a directed acyclic graph with several nodes inside. NASNet searches for the operation types and the topological connections in the cell and repeat the searched cell to form the whole network architecture. The depth of the architecture (i.e., the number of repetitions of the cell), the widths and the occurrences of down-sampling operations are all manually set. Afterwards, many works adopt a similar cell-based search space. However, architectures generated by cell-based search spaces are not friendly in terms of latency or FLOPs. Then MnasNet stacks MBConvs defined in MobileNetV2 to construct a search space for searching efficient architectures. Some works simplify the search space by searching for the expansion ratios and kernel sizes of MBConv layers.

Some works study more about the search space. Liu et al. proposes a hierarchical search space that allows flexible network topologies (directed acyclic graphs) at each level of the hierarchies. Auto-DeepLab creatively designs a two-level hierarchical search space for semantic segmentation networks. CAS customizes the search space design for real-time segmentation networks. RandWire explores randomly wired architectures by designing network generators that produce new families of models for searching. Our proposed method designs a densely connected search space beyond conventional search constrains to generate the architecture with a better trade-off between accuracy and model cost.

Some early works propose to search architectures based on reinforcement learning (RL) methods. Then evolutionary algorithm (EA) based methods achieve great performance. However, RL and EA based methods bear huge computation cost. As a result, ENAS proposes to use weight sharing for reducing the search cost.

Recently, the emergence of differentiable NAS methods and one-shot methods greatly reduces the search cost and achieves superior results. DARTS is the first work to utilize the gradient-based method to search neural architectures. They relax the architecture representation as a super network by assigning continuous weights to the candidate operations. They first search on a small dataset, e.g., CIFAR-10 , and then apply the architecture to a large dataset, e.g., ImageNet , with some manual adjustments. ProxylessNAS reduces the memory consumption by adopting a dropping path strategy and conducts search directly on the large scale dataset, i.e., ImageNet. FBNet searches on the subset of ImageNet and uses the Gumbel Softmax function to better optimize the distribution of architecture probabilities. TAS utilizes a differentiable NAS scheme to search and prune the width and depth of the network and uses knowledge distillation (KD) to promote the performance of the pruned network. FNA proposes to adapt the neural network to new tasks with low cost by a parameter remapping mechanism and differentiable NAS. It is challenging for differentiable/one-shot NAS methods to search for more flexible architectures as they need to integrate all sub-architectures into the super network. The proposed DenseNAS tends to solve this problem by integrating a densely connected search space into the differentiable paradigm and explores more flexible search schemes in the network.

Method

In this section, we first introduce how to design the search space targeted at a more flexible search. A routing block is proposed to construct the densely connected super network. Secondly, we describe the method of relaxing the search space into a continuous representation. Then, we propose a chained cost estimation algorithm to approximate the model cost during the search. Finally, we describe the whole search procedure.

As shown in Fig. 2, we define our search space using the following three terms, i.e., (basic layer, routing block and dense super Network). Firstly, a basic layer is defined as a set of all the candidate operations. Then we propose a novel routing block which can aggregate tensors from different routing blocks and transmit tensors to multiple other routing blocks. Finally, the search space is constructed as a dense super network with many routing blocks where there are various paths to transmit tensors.

We define the basic layer to be the elementary structure in our search space. One basic layer represents a set of candidate operations which include MBConvs and the skip connection. MBConvs are with kernel sizes of $\{3,5,7\}$ and expansion ratios of $\{3,6\}$. The skip connection is for the depth search. If the skip connection is chosen, the corresponding layer is removed from the resulting architecture.

1.2 Routing Block

For the purpose of establishing various paths in the super network, we propose the routing block with the ability of aggregating tensors from preceding routing blocks and transmit tensors to subsequent ones. We divide the routing block into two parts, shape-alignment layers and basic layers.

Shape-alignment layers exist in the form of several parallel branches, while every branch is a set of candidate operations. They take input tensors with different shapes (including widths and spatial resolutions) which come from multiple preceding routing blocks and transform them into tensors with the same shape. As shape-alignment layers are required for all routing blocks, we exclude the skip connection in candidate operations of them. Then tensors processed by shape-alignment layers are aggregated and sent to several basic layers. The subsequent basic layers are used for feature extraction whose depth can also be searched.

1.3 Dense Super Network

Many previous works manually set a fixed number of blocks, and retain all the blocks for the final architecture. Benefiting from the aforementioned structures of routing blocks, we introduce more routing blocks with various widths to construct the dense super network which is the representation of the search space. The final searched architecture is allowed to select a subset of the routing blocks and discard the others, giving the search algorithm more room.

We define the super network as $\mathcal{N}_{sup}$ and assume it to consist of $N$ routing blocks, $\mathcal{N}_{sup}=\{B_{1},B_{2},...,B_{N}\}$. The network structure is shown in Fig. 2. We partition the entire network into several stages. As Sec. 1 defines, each stage contains routing blocks with various widths and the same spatial resolution. From the beginning to the end of the super network, the widths of routing blocks grow gradually. In the early stage of the network, we set a small growing stride for the width because large width settings in the early network stage will cause huge computational cost. The growing stride becomes larger in the later stages. This design principle of the super network allows more possibilities of block counts and block widths.

We assume that each routing block in the super network connects to $M$ subsequent ones. We define the connection between the routing block $B_{i}$ and its subsequent routing block $B_{j}$ ($j>i$) as $C_{ij}$. The spatial resolutions of $B_{i}$ and $B_{j}$ are $H_{i}\times W_{i}$ and $H_{j}\times W_{j}$ respectively (normally $H_{i}=W_{i}$ and $H_{j}=W_{j}$). We set some constraints on the connections to avoid the stride of the spatial down-sampling exceeding $2$. Specifically, $C_{ij}$ only exists when $j-i\leq M$ and $H_{i}/H_{j}\leq 2$. Following the above paradigms, the search space is constructed as a dense super network based on the connected routing blocks.

2 Relaxation of Search Space

We integrate our search space by relaxing the architectures into continuous representations. The relaxation is implemented on both the basic layer and the routing block. We can search for architectures via back-propagation in the relaxed search space.

2.2 Relaxation in the Routing Block

We assume that the routing block $B_{i}$ outputs the tensor $b_{i}$ and connects to $m$ subsequent blocks. To relax the block connections as a continuous representation, we assign each output path of the block an architecture parameter. Namely the path from $B_{i}$ to $B_{j}$ has a parameter $\beta_{ij}$. Similar to how we compute the architecture weight of each operation above, we compute the probability of each path using a softmax function over all paths between the two routing blocks:

For routing block $B_{i}$, we assume it takes input tensors from its $m^{\prime}$ preceding routing blocks ($B_{i-m^{\prime}}$, $B_{i-m^{\prime}+1}$, $B_{i-m^{\prime}+2}$ … $B_{i-1}$). As shown in Fig. 2, the input tensors from these routing blocks differ in terms of width and spatial resolution. Each input tensor is transformed to a same size by the corresponding branch of shape-alignment layers in $B_{i}$. Let $H_{ik}$ denotes the $k$th transformation branch in $B_{i}$ which is applied to the input tensor from $B_{i-k}$, where $k=1\dots m^{\prime}$. Then the input tensors processed by shape-alignment layers are aggregated by a weighted-sum using the path probabilities,

It is worth noting that the path probabilities are normalized on the output dimension but applied on the input dimension (more specifically on the branches of shape-alignment layers). One of the shape-alignment layers is essentially a weighted-sum mixture of the candidate operations. The layer-level parameters $\alpha$ control which operation to be selected, while the outer block-level parameters $\beta$ determine how blocks connect.

3 Chained Cost Estimation Algorithm

We propose to optimize both the accuracy and the cost (latency/FLOPs) of the model. To this end, the model cost needs to be estimated during the search. In conventional cascaded search spaces, the total cost of the whole network can be computed as a sum of all the blocks. Instead, the global effects of connections on the predicted cost need to be taken into consideration in our densely connected search space. We propose a chained cost estimation algorithm to better approximate the model cost.

We create a lookup table which records the cost of each operation in the search space. The cost of every operation is measured separately. During the search, the cost of one basic layer is estimated as follows,

We design a loss function with the cost-based regularization to achieve the multi-objective optimization:

where $\lambda$ and $\tau$ are the hyper-parameters to control the magnitude of the model cost term.

4 Search Procedure

Benefiting from the continuously relaxed representation of the search space, we can search for the architecture by updating the architecture parameters (introduced in Sec. 3.2) using stochastic gradient descent. We find that at the beginning of the search process, all the weights of the operations are under-trained. The operations or architectures which converge faster are more likely to be strengthened, which leads to shallow architectures. To tackle this, we split our search procedure into two stages. In the first stage, we only optimize the weights for enough epochs to get operations sufficiently trained until the accuracy of the model is not too low. In the second stage, we activate the architecture optimization. We alternatively optimize the operation weights by descending $\nabla_{w}\mathcal{L}_{train}(w,\alpha,\beta)$ on the training set, and optimize the architecture parameters by descending $\nabla_{\alpha,\beta}\mathcal{L}_{val}(w,\alpha,\beta)$ on the validation set. Moreover, a dropping-path training strategy is adopted to decrease memory consumption and decouple different architectures in the super network.

Experiments

In this section, we first show the performance with the MobileNetV2 -based search space on ImageNet classification. Then we apply the architectures searched on ImageNet to object detection on COCO . We further extend our DenseNAS to the ResNet -based search space. Finally, we conduct some ablation studies and analysis. The implementation details are provided in the appendix.

We implement DenseNAS on the MobileNetV2 -based search space, set the GPU latency as our secondary optimization objective, and search models with different sizes under multiple latency optimization magnitudes (defined in Eq. 8). The ImageNet results are shown in Tab. 1. We divide Tab. 1 into several parts and compare DenseNAS models with both manually designed models and NAS models. DenseNAS achieves higher accuracies with both fewer FLOPs and lower latencies. Note that for FBNet-A, the group convolution in the 1$\times$1 conv and the channel shuffle operation are used, which do not exist in FBNet-B, -C and Proxyless. In the compared NAS methods , the block counts and block widths in the search space are set and adjusted manually. DenseNAS allocates block counts and block widths automatically. We further visualize the results in Fig. 3, which clearly demonstrates that DenseNAS achieves a better trade-off between accuracy and latency. The searched architectures are shown in Fig. 4.

2 Generalization Ability on COCO Object Detection

We apply the searched DenseNAS networks on the COCO object detection task to evaluate the generalization ability of DenseNAS networks and show the results in Tab. 2. We choose two commonly used object detection frameworks RetinaNet and SSDLite to conduct our experiments. All the architectures shown in Tab. 2 are utilized as the backbone networks in the detection frameworks. The experiments are performed based on the MMDetection framework.

We compare our results with both manually designed and NAS models. Results of MobileNetV2 , FBNet and ProxylessNAS are obtained by our re-implementation and all models are trained under the same settings and hyper-parameters for fair comparisons. DetNAS is a recent work that aims at searching the backbone architectures directly on object detection. Though DenseNAS searches on the ImageNet classification task and applies the searched architectures on detection tasks, our DenseNAS models still obtain superior detection performance in terms of both accuracy and FLOPs. The superiority over the compared methods demonstrates the great generalization ability of DenseNAS networks.

3 Performance on ResNet-based Search Space

We apply our DenseNAS framework on the ResNet -based search space to further evaluate the generalization ability of our method. It is convenient to implement DenseNAS on ResNet as we set the candidate operations in the basic layer as the basic block defined in ResNet and the skip connection. The ResNet-based search space is also constructed as a densely connected super network.

We search for several architectures with different FLOPs and compare them with the original ResNet models on the ImageNet classification task in Tab. 3. We further replace all the basic blocks in DenseNAS-R2 with the bottleneck blocks and obtain DenseNAS-R3 to compare with ResNet-50-B and the NAS model RandWire-WS, C=109 (WS, C=109). Though WS, C=109 achieves a higher accuracy, the FLOPs increases 600M, which is a great number, 17.6% of DenseNAS-R3. Besides, WS C=109 uses separable convolutions which greatly decrease the FLOPs while DenseNAS-R3 only contains plain convolutions. Moreover, RandWire networks are unfriendly to inference on existing hardware for the complicated connection patterns. Our proposed DenseNAS promotes the accuracy of ResNet-18, -34 and -50-B by 1.5%, 0.5% and 0.3% with 200M, 600M, 680M fewer FLOPs and 1.5ms, 2.4ms, 6.1ms lower latency respectively. We visualize the comparison results in Fig. 5 and the performance on the ResNet-based search space further demonstrates the great generalization ability and effectiveness of DenseNAS.

4 Ablation Study and Analysis

To further demonstrate the effectiveness of our proposed densely connected search space, we conduct the same search algorithm used in DenseNAS on the search spaces of FBNet and ProxylessNAS as well as a new search space which is constructed following the settings of block counts and block widths in MobileNetV2. The three search spaces are denoted as FBNet-SS, Proxyless-SS and MBV2-SS respectively. All the search/training settings and hyper-parameters are the same as that we use for DenseNAS. The results are shown in Tab. 4 and DenseNAS achieves the highest accuracy with the lowest latency.

As random search is treated as an important baseline to validate NAS methods. We conduct random search experiments and show the results in Tab. 1. We randomly sample 15 models in our search space whose FLOPs are similar to DenseNAS-C. Then we train every model for 5 epochs on ImageNet. Finally, we select the one with the highest validation accuracy and train it under the same settings as DenseNAS. The total search cost of the random search is the same as DenseNAS. We observe that DenseNAS-C is 1% accuracy higher compared with the randomly searched model, which proves the effectiveness of DenseNAS.

We explore the effect of the maximum number of connections between routing blocks in the search space. We set the maximum connection number as 4 in DenseNAS. Then we try more options and show the results in Tab. 5. When we set the connection number to 3, the searched model gets worse performance. We attribute this to the search space shrinkage which causes the loss of many possible architectures with good performance. As we set the number to 5 and the search process takes the same number of epochs as DenseNAS, i.e. 150 epochs. The performance of the searched model is not good, even worse than that of the connection number 3. Then we increase the search epochs to 200 and the search process achieves a comparable result with DenseNAS. This phenomenon indicates that larger search spaces need more search cost to achieve comparable/better results with/than smaller search spaces with some added constraints.

As the super network is densely connected and the final architecture is derived based on the total transition probability, the model cost estimation needs to take the effects of all the path probabilities on the whole network into consideration. We try a local cost estimation strategy that does not involve the global connection effects on the whole super network. Specifically, we compute the cost of the whole network by summing the cost of every routing block during the search as follows, while the transition probability $p_{ji}$ is only used for computing the cost of each individual block rather than the whole network.

where all definitions in the equation are the same as that in Eq. 6. We randomly generate the architecture parameters ($\alpha$ and $\beta$) to derive the architectures. Then we draw the approximated cost values computed by local cost estimation and our proposed chained cost estimation respectively, and compare with the real cost values in Fig. 6. In this experiment, we take FLOPs as the model cost because the FLOPs is easier to measure than latency. 1,500 models are sampled in total. The results show that the predicted cost values computed by our chained cost estimation algorithm has a much stronger correlation with the real values and approximate more to the real ones. As the predicted values are computed based on the randomly generated architecture parameters which are not binary parameters, there are still differences between the predicted and real values.

We visualize the searched architectures in Fig. 4. It shows that DenseNAS-B and -C have one more block in the last stage than other architectures, which indicates enlarging the depth in the last stage of the network tends to obtain a better accuracy. Moreover, the smallest architecture DenseNAS-A whose FLOPs is only 251M has one fewer block than DenseNAS-B and -C to decrease the model cost. The structures of the final searched architectures show the great flexibility of DenseNAS.

Conclusion

We propose a densely connected search space for more flexible architecture search, DenseNAS. We tackle the limitations in previous search space design in terms of the block counts and widths. The novelly designed routing blocks are utilized to construct the search space. The proposed chained cost estimation algorithm aims at optimizing both accuracy and model cost. The effectiveness of DenseNAS is demonstrated on both MobileNetV2- and ResNet- based search spaces. We leave more applications, e.g. semantic segmentation, face detection, pose estimation, and more network-based search space implementations, e.g. MobileNetV3 , ShuffleNet and VarGNet , for future work.

Acknowledgement

This work was supported by National Key R&D Program of China (No. 2018YFB1402600), National Natural Science Foundation of China (NSFC) (No. 61876212, No. 61733007 and No. 61572207), and HUST-Horizon Computer Vision Research Center. We thank Liangchen Song, Kangjian Peng and Yingqing Rao for the discussion and assistance.

References

Appendix A Appendix

Before the search process, we build a lookup table for every operation latency of the super network as described in Sec. 3.3. We set the input shape as $(3,224,224)$ with the batch size of $32$ and measure each operation latency on one TITAN-XP GPU. All models and experiments are implemented using PyTorch .

For the search process, we randomly choose $100$ classes from the original 1K-class ImageNet training set. We sample $20\%$ data of each class from the above subset as the validation set. The original validation set of ImageNet is only used for evaluating our final searched architecture. The search process takes 150 epochs in total. We first train the operation weights for 50 epochs on the divided training set. For the last 100 epochs, the updating of architecture parameters ($\alpha,\beta$) and operation weights ($w$) alternates in each epoch. We use the standard GoogleNet data augmentation for the training data preprocessing. We set the batch size to $352$ on $4$ Tesla V100 GPUs. The SGD optimizer is used with $0.9$ momentum and $4\times 10^{-5}$ weight decay to update the operation weights. The learning rate decays from $0.2$ to $1\times 10^{-4}$ with the cosine annealing schedule . We use the Adam optimizer with $10^{-3}$ weight decay, $\beta=(0.5,0.999)$ and a fixed learning rate of $3\times 10^{-4}$ to update the architecture parameters.

For retraining the final derived architecture, we use the same data augmentation strategy as the search process on the whole ImageNet dataset. We train the model for $240$ epochs with a batch size of $1024$ on $8$ TITAN-XP GPUs. The optimizer is SGD with $0.9$ momentum and $4\times 10^{-5}$ weight decay. The learning rate decays from 0.5 to $1\times 10^{-4}$ with the cosine annealing schedule.

A.2 Viterbi Algorithm for Block Deriving

The Viterbi Algorithm is widely used in dynamic programming which targets at finding the most likely path between hidden states. In DenseNAS, only a part of routing blocks in the super network are retained to construct the final architecture. As described in Sec. 3.4, we implement the Viterbi algorithm to derive the final sequence of blocks. We treat the routing block in the super network as each hidden state in the Viterbi algorithm. The path probability $p_{ij}$ serves as the transition probability from routing block $B_{i}$ to $B_{j}$. The total algorithm is described in Algo. 1. The derived block sequence holds the maximum transition probability.

A.3 Dropping-path Search Strategy

A.4 Implementation Details of ResNet Search

We design the ResNet-based search space as follows. As enlarging the kernel size of the ResNet block causes a huge computation cost increase, the candidate operations in the basic layer only include the basic block and the skip connection. That means we aim at width and depth search for ResNet networks. During the search, the batch size is set as 512 on 4 Tesla V100 GPUs. The search process takes 70 epochs in total and we start to update the architecture parameters from epoch 10. We set all the other search settings and hyper-parameters the same as that in the MobileNetV2 search. For the architecture retraining, the same training settings and hyper-parameters are used as that for architectures searched in the MobileNetV2-based search space. The architectures searched by DenseNAS are shown in Tab. A.1.

A.5 Experimental Comparison of Cost Estimation Method

We study the design of the model cost estimation algorithm in Sec. 4.5. $1,500$ models are derived based on the randomly generated architecture parameters. Cost values predicted by our proposed chained cost estimation algorithm demonstrate a stronger correlation with the real values and more accurate prediction results than the compared local cost estimation strategy. We further perform the same search process as DenseNAS on the MobileNetV2 -based search space with the local estimation strategy and show the searched results in Tab. 7. DenseNAS with the chained cost estimation algorithm shows a higher accuracy with lower latency and fewer FLOPs. It proves the effectiveness of the chained cost estimation algorithm on achieving a good trade-off between accuracy and model cost.