Improve Ranking Correlation of Super-net through Training Scheme from One-shot NAS to Few-shot NAS

Different from CVPR2021 Lightweight NAS Challenge Track1, which only searches the number of channels, this track in 2022 searches for the number of layers of the network and the number of channels at each layer. It would be very difficult to construct a super-net like what SPOS[4] and FairNAS[2] have done, so we construct that according to the ways of OFA[1] and BigNAS[9] with weight entanglement. The super-net is constructed based on resnet48, which contains 4 stages. The number of searchable channels is the product of the basic channel and scaling ratio. The scaling ratio in all stages is [1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7]. The number of searchable blocks is from 2 to 8. The details of search space in every stage are shown in Table 1.

For the one-shot super-net, we follow the pipeline of BigNAS[9]. In detail, we first construct the super-net based on the search space of $S$ with weights $\boldsymbol{\theta}$, and then we train the largest sub-network $S_{max}$ with the cross-entropy loss $\mathcal{L}_{CE}$. Finally, we distill the smallest sub-network $S_{min}$ and two sub-networks sampled randomly respectively with the Kullback–Leibler (KL) loss (distance) $\mathcal{L}_{KL}$ between the logit of them and the one of $S_{max}$. Algorithm 1 shows the training process.

As for the few-shot super-net, to alleviate the interference among the subnets whose weights are shared, the proposed method follows the way of few-shot NAS, which disentangles the shared weights by splitting the super-net into multi-subnet during the training process. In this section, we introduce the split way for the resnet-like search space.

There are two kinds of channel number of convolution layers that need to be searched in one block, in which the second convolution layer controls the channel number of the output features. Because of the existence of identity mapping in the block of the resnet, the channel number of the output features is the same in every block in the same stage. In other words, the second channel number in different blocks in the same stage is fixed in one search process. Based on this analysis, we divide the search space following the dimensionality of the number in the second convolution layer. Specifically, the candidates for the second channel number are 88, 96, 104, 112, 120, and 128 for the second stage of this search space, whose basic channel number is 128. Instead of retaining only one group based on the vanilla one-shot NAS, we extend the number of the group up to the number of the candidates, that is to say, we divide the second stage into six groups, each of which contains the same architecture except for the channel number of the second convolution layer. We apply this split way to all stages except for the stem stage. The left of Figure 1 shows the details of the architectures of all stages.

According to the aforementioned ways, we first train the super-net with one group by Algorithm 1 until convergence (the same as the one-shot way), and then the super-net with two groups is trained following the same algorithm with the weights inheriting from the ones of the super-net with one group (The details of weights inheriting can be seen in the right of Figure 1). We repeat the training process until the number of the group increases to maximum.

We firstly train the one-shot super-net for 100 epochs with batch size 128 and apply a stochastic gradient descent optimizer with the momentum of 0.9. The learning rate (LR) is decreased from an initial value of 0.1 to 0 with a linear-warmup cosine learning rate decay strategy, and the warmup epoch is 5. For the largest subnet, the weights are regularized with a weight decay of 0.00001, and the dropout probability is 0.2. The data augmentation includes brightness and contrast changes with a probability of 0.5, random rotation within 15 degrees, and random clipping of 224*224.

The training configurations of few-shot NAS are mostly the same as that of one-shot NAS, except that epoch is set to 20 and there is no warm-up for the learning rate. In addition, different network modules are assigned different learning rates according to their probability of being sampled, to ensure the same order of magnitude in the training times. As shown in Table 2, taking group 3 as an example, the LR of the module with channel number 112 in stage 1 is 4 times of the basic LR. LR of the module with channel number 232 in stage 2 and channel number 464 in stage 3 is 5 times of basic LR, and the learning rate of the rest modules is the basic learning rate.

We selected the corresponding model at each training super-net stage to calculate the score, as shown in Table 3.

As [11] said, the ranking does not increase steadily during training as accuracy does. Therefore, it is important to select the most appropriate epoch for the super-net. However, the large test set and a large number of subnets to be evaluated make it time-consuming to calculate rankings on each epoch. To solve this problem, we standalone trained 30 subnets with the part of valid datasets (the amount is $\frac{1}{20}$ of the original datasets’) to obtain the target accuracy, and then used the same datasets to evaluate the accuracy of the 30 subnets inherited from the weight of the current super-net after each epoch when training the super-net. Therefore, we obtained the change curve of Kendall tau with the epoch as shown in Figure 2, which shows that the ranking is rising but very unstable. The super-net with late epoch and large Kendall tau is selected as the measured model.

We discuss the different ranking performances of super-nets based on different architecture sampling strategies in this paragraph. To reduce the test time, we evaluate the super-net trained with the one-shot method on the part of datasets mentioned before. Table 4 shows the results of the uniform sample and fair sample, in which the former is better than the latter obviously, opposite to the results of [2]. This may be caused by the wrong usage of the fair samples. In [2], it requires the same amount of candidates in the search space, which is not satisfied strictly in this search space.