NAS-FCOS: Fast Neural Architecture Search for Object Detection

Introduction

Object detection is one of the fundamental tasks in computer vision, and has been researched extensively. In the past few years, state-of-the-art methods for this task are based on deep convolutional neural networks (such as Faster R-CNN , RetinaNet ), due to their impressive performance. Typically, the designs of object detection networks are much more complex than those for image classification, because the former need to localize and classify multiple objects in an image simultaneously while the latter only need to output image-level labels. Due to its complex structure and numerous hyper-parameters, designing effective object detection networks is more challenging and usually needs much manual effort.

On the other hand, Neural Architecture Search (NAS) approaches have been showing impressive results on automatically discovering top-performing neural network architectures in large-scale search spaces. Compared to manual designs, NAS methods are data-driven instead of experience-driven, and hence need much less human intervention. As defined in , the workflow of NAS can be divided into the following three processes: $1$) sampling architecture from a search space following some search strategies; $2$) evaluating the performance of the sampled architecture; and $3$) updating the parameters based on the performance.

One of the main problems prohibiting NAS from being used in more realistic applications is its search efficiency. The evaluation process is the most time consuming part because it involves a full training procedure of a neural network. To reduce the evaluation time, in practice a proxy task is often used as a lower cost substitution. In the proxy task, the input, network parameters and training iterations are often scaled down to speedup the evaluation. However, there is often a performance gap for samples between the proxy tasks and target tasks, which makes the evaluation process biased. How to design proxy tasks that are both accurate and efficient for specific problems is a challenging problem. Another solution to improve search efficiency is constructing a supernet that covers the complete search space and training candidate architectures with shared parameters . However, this solution leads to significantly increased memory consumption and restricts itself to small-to-moderate sized search spaces.

To our knowledge, studies on efficient and accurate NAS approaches to object detection networks are rarely touched, despite its significant importance. To this end, we present a fast and memory saving NAS method for object detection networks, which is capable of discovering top-performing architectures within significantly reduced search time. Our overall detection architecture is based on FCOS , a simple anchor-free one-stage object detection framework, in which the feature pyramid network and prediction head are searched using our proposed NAS method.

Our main contributions are summarized as follows.

In this work, we propose a fast and memory-efficient NAS method for searching both FPN and head architectures, with carefully designed proxy tasks, search space and evaluation strategies, which is able to find top-performing architectures over $3,000$ architectures using $28$ GPU-days only.

Specifically, this high efficiency is enabled with the following designs.

$-$ Developing a fast proxy task training scheme by skipping the backbone finetuning stage;

$-$ Adapting progressive search strategy to reduce time cost taken by the extended search space;

$-$ Using a more discriminative criterion for evaluation of searched architectures.

$-$ Employing an efficient anchor-free one-stage detection framework with simple post processing;

Using NAS, we explore the workload relationship between FPN and head, proving the importance of weight sharing in head.

We show that the overall structure of NAS-FCOS is general and flexible in that it can be equipped with various backbones including MobileNetV$2$, ResNet-$50$, ResNet-$101$ and ResNeXt-$101$, and surpasses state-of-the-art object detection algorithms using comparable computation complexity and memory footprint. More specifically, our model can improve the AP by $1.5\sim 3.5$ points on all above models comparing to their FCOS counterparts.

Related Work

The frameworks of deep neural networks for object detection can be roughly categorized into two types: one-stage detectors and two-stage detectors .

Two-stage detection frameworks first generate class-independent region proposals using a region proposal network (RPN), and then classify and refine them using extra detection heads. In spite of achieving top performance, the two-stage methods have noticeable drawbacks: they are computationally expensive and have many hyper-parameters that need to be tuned to fit a specific dataset.

In comparison, the structures of one-stage detectors are much simpler. They directly predict object categories and bounding boxes at each location of feature maps generated by a single CNN backbone.

Note that most state-of-the-art object detectors (including both one-stage detectors and two-stage detectors ) make predictions based on anchor boxes of different scales and aspect ratios at each convolutional feature map location. However, the usage of anchor boxes may lead to high imbalance between object and non-object examples and introduce extra hyper-parameters. More recently, anchor-free one-stage detectors have attracted increasing research interests, due to their simple fully convolutional architectures and reduced consumption of computational resources.

2 Neural Architecture Search

NAS is usually time consuming. We have seen great improvements from $24,000$ GPU-days to $0.2$ GPU-day . The trick is to first construct a supernet containing the complete search space and train the candidates all at once with bi-level optimization and efficient weight sharing . But the large memory allocation and difficulties in approximated optimization prohibit the search for more complex structures.

Recently researchers propose to apply single-path training to reduce the bias introduced by approximation and model simplification of the supernet. DetNAS follows this idea to search for an efficient object detection architecture. One limitation of the single-path approach is that the search space is restricted to a sequential structure. Single-path sampling and straight through estimate of the weight gradients introduce large variance to the optimization process and prohibit the search for more complex structures under this framework. Within this very simple search space, NAS algorithms can only make trivial decisions like kernel sizes for manually designed modules.

Object detection models are different from single-path image classification networks in their way of merging multi-level features and distributing the task to parallel prediction heads. Feature pyramid networks (FPNs) , designed to handle this job, plays an important role in modern object detection models. NAS-FPN targets on searching for an FPN alternative based on one-stage framework RetinaNet . Feature pyramid architectures are sampled with a recurrent neural network (RNN) controller. The RNN controller is trained with reinforcement learning (RL). However, the search is very time-consuming even though a proxy task with ResNet-10 backbone is trained to evaluate each architecture.

Since all these three kinds of research ( and ours) focus on object detection framework, we demonstrate the differences among them that DetNAS aims to search for the designs of better backbones, while NAS-FPN searches the FPN structure, and our search space contains both FPN and head structure.

To speed up reward evaluation of RL-based NAS, the work of proposes to use progressive tasks and other training acceleration methods. By caching the encoder features, they are able to train semantic segmentation decoders with very large batch sizes very efficiently. In the sequel of this paper, we refer to this technique as fast decoder adaptation. However, directly applying this technique to object detection tasks does not enjoy similar speed boost, because they are either not in using a fully-convolutional model or require complicated post processing that are not scalable with the batch size .

To reduce the post processing overhead, we resort to a recently introduced anchor-free one-stage framework, namely, FCOS , which significantly improve the search efficiency by cancelling the processing time of anchor-box matching in RetinaNet.

Compared to its anchor-based counterpart, FCOS significantly reduces the training memory footprint while being able to improve the performance.

Our Approach

In our work, we search for anchor-free fully convolutional detection models with fast decoder adaptation. Thus, NAS methods can be easily applied.

We base our search algorithm upon a one-stage framework FCOS due to its simplicity. Our training tuples $\{(\mathbf{x},Y)\}$ consist of input image tensors $\mathbf{x}$ of size $(3\times H\times W)$ and FCOS output targets $Y$ in a pyramid representation, which is a list of tensors $\mathbf{y}_{l}$ each of size $((K+4+1)\times H_{l}\times W_{l})$ where $H_{l}\times W_{l}$ is feature map size on level $p$ of the pyramid. $(K+4+1)$ is the output channels of FCOS, the three terms are length-$K$ one-hot classification labels, $4$ bounding box regression targets and $1$ centerness factor respectively.

The network $g:\mathbf{x}\rightarrow\hat{Y}$ in original FCOS consists of three parts, a backbone $b$, FPN $f$ and multi-level subnets we call prediction heads $h$ in this paper. First backbone $b:\mathbf{x}\rightarrow C$ maps the input tensor to a set of intermediate-leveled features $C=\{\mathbf{c}_{3},\mathbf{c}_{4},\mathbf{c}_{5}\}$, with resolution $(H_{i}\times W_{i})=(H/2^{i}\times W/2^{i})$. Then FPN $f:C\rightarrow P$ maps the features to a feature pyramid $P=\{\mathbf{p}_{3},\mathbf{p}_{4},\mathbf{p}_{5},\mathbf{p}_{6},\mathbf{p}_{7}\}$. Then the prediction head $h:\mathbf{p}\rightarrow\mathbf{y}$ is applied to each level of $P$ and the result is collected to create the final prediction. To avoid overfitting, same $h$ is often applied to all instances in $P$.

Since objects of different scales require different effective receptive fields, the mechanism to select and merge intermediate-leveled features $C$ is particularly important in object detection network design. Thus, most researches are carried out on designing $f$ and $h$ while using widely-adopted backbone structures such as ResNet . Following this principle, our search goal is to decide when to choose which features from $C$ and how to merge them.

To improve the efficiency, we reuse the parameters in $b$ pretrained on target dataset and search for the optimal structures after that. For the convenience of the following statement, we call the network components to search for, namely $f$ and $h$, together the decoder structure for the objection detection network.

$f$ and $h$ take care of different parts of the detection job. $f$ extracts features targeting different object scales in the pyramid representations $P$, while $h$ is a unified mapping applied to each feature in $P$ to avoid overfitting. In practice, people seldom discuss the possibility of using a more diversified $f$ to extract features at different levels or how many layers in $h$ need to be shared across the levels. In this work, we use NAS as an automatic method to test these possibilities.

2 Search Space

Considering the different functions of $f$ and $h$, we apply two search space respectively. Given the particularity of FPN structure, a basic block with new overall connection and $f$’s output design is built for it. For simplicity, sequential space is applied for $h$ part.

We replace the cell structure with atomic operations to provide even more flexibility. To construct one basic block, we first choose two layers $\mathbf{x}_{1}$, $\mathbf{x}_{2}$ from the sampling pool $X$ at id1, id2, then two operations op1, op2 are applied to each of them and an aggregation operation agg merges the two output into one feature. To build a deep decoder structure, we apply multiple basic blocks with their outputs added to the sampling pool. Our basic block $bb_{t}:X_{t-1}\rightarrow X_{t}$ at time step $t$ transforms the sampling pool $X_{t-1}$ to $X_{t}=X_{t-1}\cup\{\mathbf{x}_{t}\}$, where $\mathbf{x}_{t}$ is the output of $bb_{t}$.

The candidate operations are listed in Table 1. We include only separable/depth-wise convolutions so that the decoder can be efficient. In order to enable the decoder to apply convolutional filters on irregular grids, here we have also included deformable $3\times 3$ convolutions . For the aggregation operations, we include element-wise sum and concatenation followed by a $1\times 1$ convolution.

The decoder configuration can be represented by a sequence with three components, FPN configuration, head configuration and weight sharing stages. We provide detailed descriptions to each of them in the following sections. The complete diagram of our decoder structure is shown in Fig. 1.

As mentioned above, the FPN $f$ maps the convolutional features $C$ to $P$. First, we initialize the sampling pool as $X_{0}=C$. Our FPN is defined by applying the basic block $7$ times to the sampling pool, $f:=bb_{1}^{f}\circ bb_{2}^{f}\circ\cdots\circ bb_{7}^{f}$. To yield pyramid features $P$, we collect the last three basic block outputs $\{\mathbf{x}_{5},\mathbf{x}_{6},\mathbf{x}_{7}\}$ as $\{\mathbf{p}_{3},\mathbf{p}_{4},\mathbf{p}_{5}\}$.

To allow shared information across all layers, we use a simple rule to create global features. If there is some dangling layer $\mathbf{x}_{t}$ which is not sampled by later blocks $\{bb_{i}^{f}|i>t\}$ nor belongs to the last three layers $t<5$, we use element-wise add to merge it to all output features

Same as the aggregation operations, if the features have different resolution, the smaller one is upsampled with bilinear interpolation.

To be consistent with FCOS, $\mathbf{p}_{6}$ and $\mathbf{p}_{7}$ are obtained via a $3\times 3$ stride-$2$ convolution on $\mathbf{p}_{5}$ and $\mathbf{p}_{6}$ respectively.

2.2 Prediction Head Search Space

Prediction head $h$ maps each feature in the pyramid $P$ to the output of corresponding $\mathbf{y}$, which in FCOS and RetinaNet, consists of four $3\times 3$ convolutions. To explore the potential of the head, we therefore extend a sequential search space for its generation. Specifically, our head is defined as a sequence of six basic operations. Compared with candidate operations in the FPN structures, the head search space has two slight differences. First, we add standard convolution modules (including conv$1$x$1$ and conv$3$x$3$) to the head sampling pool for better comparison. Second, we follow the design of FCOS by replacing all the Batch Normalization (BN) layers to Group Normalization (GN) in the operations sampling pool of head, considering that head needs to share weights between different levels, which causes BN invalid. The final output of head is the output of the last (sixth) layer.

2.3 Searching for Head Weight Sharing

To add even more flexibility and understand the effect of weight sharing in prediction heads, we further add an index $i$ as the location where the prediction head starts to share weights. For every layer before stage $i$, the head $h$ will create independent set of weights for each FPN output level, otherwise, it will use the global weights for sharing purpose.

Considering the independent part of the heads being extended FPN branch and the shared part as head with adaptive-length, we can further balance the workload for each individual FPN branch to extract level-specific features and the prediction head shared across all levels.

3 Search Strategy

RL based strategy is applied to the search process. We rely on an LSTM-based controller to predict the full configuration. We consider using a progressive search strategy rather than the joint search for both FPN structure and prediction head part, since the former requires less computing resources and time cost than the latter. The training dataset is randomly split into a meta-train $D_{t}$ and meta-val $D_{v}$ subset. To speed up the training, we fix the backbone network and cache the pre-computed backbone output $C$. This makes our single architecture training cost independent from the depth of backbone network. Taking this advantage, we can apply much more complex backbone structures and utilize high quality multilevel features as our decoder’s input. We find that the process of backbone finetuning can be skipped if the cached features are powerful enough. Speedup techniques such as Polyak weight averaging are also applied during the training.

The most widely used detection metric is average precision (AP). However, due to the difficulty of object detection task, at the early stages, AP is too low to tell the good architectures from the bad ones, which makes the controller take much more time to converge. To make the architecture evaluation process easier even at the early stages of the training, we therefore use negative loss sum as the reward instead of average precision:

where $L_{cls}$, $L_{reg}$, $L_{ctr}$ are the three loss terms in FCOS. Gradient of the controller is estimated via proximal policy optimization (PPO) .

Experiments

We design a fast proxy task for evaluating the decoder architectures sampled in the searching phase. PASCAL VOC is selected as the proxy dataset, which contains $5715$ training images with object bounding box annotations of $20$ classes. Transfer capacity of the structures can be illustrated since the search and full training phase use different datasets. The VOC training set is randomly split into a meta-train set with $4,000$ images and a meta-val set with $1715$ images. For each sampled architecture, we train it on meta-train and compute the reward (2) on meta-val. Input images are resized to short size $384$ and then randomly cropped to $384\times 384$. Target object sizes of interest are scaled correspondingly. We use Adam optimizer with learning rate $8$e$-4$ and batch size $200$. Polyak averaging is applied with the decay rates of $0.9$. The decoder is evaluated after $300$ iterations. As we use fast decoder adaptation, the backbone features are fixed and cached during the search phase. To enhance the cached backbone features, we first initialize them with pre-trained weights provided by open-source implementation of FCOShttps://tinyurl.com/FCOSv1 and then finetune on VOC using the training strategies of FCOS. Note that the above finetuning process is only performed once at the begining of the search phase.

A progressive strategy is used for the search of $f$ and $h$. We first search for the FPN part and retain the original head. All operations in the FPN structure have $64$ output channels. The decoder inputs $C$ are resized to fit output channel width of FPN via $1\times 1$ convolutions. After this step, a searched FPN structure is fixed and the second stage searching for the head will be started based on it. Most parameters for searching head are identical to those for searching FPN structure, with the exception that the output channel width is adjusted from $64$ to $128$ to deliver more information.

For the FPN search part, the controller model nearly converged after searching over $2.8$K architectures on the proxy task as shown in Fig. 2. Then, the top-$20$ best performing architectures on the proxy task are selected for the next full training phase. For the head search part, we choose the best searched FPN among the top-$20$ architectures and pre-fetch its features. It takes about $600$ rounds for the controller to nearly converge, which is much faster than that for searching FPN architectures. After that, we select for full training the top-$10$ heads that achieve best performance on the proxy task. In total, the whole search phase can be finished within $4$ days using $8$ V100 GPUs.

1.2 Full Training Phase

In this phase, we fully train the searched models on the MS COCO training dataset, and select the best one by evaluating them on MS COCO validation images. Note that our training configurations are exactly the same as those in FCOS for fair comparison. Input images are resized to short size $800$ and the maximum long side is set to be $1333$. The models are trained using $4$ V100 GPUs with batch size $16$ for $90$K iterations. The initial learning rate is $0.01$ and reduces to one tenth at the $60$K-th and $80$K-th iterations. The improving tricks are applied only on the final model (w/improv).

2 Search Results

The best FPN structure is illustrated in Fig. 3. The controller identifies that deformable convolution and concatenation are the best performing operations for unary and aggregation respectively. From Fig. 4, we can see that the controller chooses to use $4$ operations (with two skip connections), rather than the maximum allowed $6$ operations. Note that the discovered “dconv + $1$x$1$ conv” structure achieves a good trade-off between accuracy and FLOPs. Compared with the original head, our searched head has fewer FLOPs/Params (FLOPs $79.24$G vs. $89.16$G, Params $3.41$M vs. $4.92$M) and significantly better performance (AP $38.7$ vs. $37.4$).

We use the searched decoder together with either light-weight backbones such as MobileNet-V2 or more powerful backbones such as ResNet-$101$ and ResNeXt-$101$ . To balance the performance and efficiency, we implement three decoders with different computation budgets: one with feature dimension of $128$ (@$128$), one with $256$ (@$256$) and another with FPN channel width $128$ and prediction head $256$ (@$128$-$256$). The results on the COCO test-dev with short side being $800$ is shown in Table 2. The searched decoder with feature dimension of $256$ (@$256$) surpasses its FCOS counterpart by $1.5$ to $3.5$ points in AP under different backbones. The one with $128$ channels (@$128$) has significantly reduced parameters and calculation, making it more suitable for resource-constrained environments. In particular, our searched model with $128$ channels and MobileNetV2 backbone suparsses the original FCOS with the same backbone by $0.8$ AP points with only $1/3$ FLOPS. The third type of decoder (@$128$-$256$) achieves a good balance between accuracy and parameters. Note that our searched model outperforms the strongest FCOS variant by $1.4$ AP points ($46.1$ vs. $44.7$) with slightly smaller FLOPs and Params. The comparison of FLOPs and number of parameters with other models are illustrated in Fig. 7 and Fig. 8 respectively.

In order to understand the importance of weight sharing in head, we add the number of layers shared by weights as an object of the search. Fig. 5 shows a trend graph of head weight sharing during search. We set $50$ structures as a statistical cycle. As the search deepens, the proportion of fully shared structures increases, indicating that on the multi-scale detection model, head weight sharing is a necessity.

We also demonstrate the comparison with other NAS methods for object detection in Table 3. Our method is able to search for twice more architectures than DetNAS per GPU-day. Note that the AP of NAS-FPN is achieved by stacking the searched FPN $7$ times, while we do not stack our searched FPN. Our model with ResNeXt-101 ($64$x$4$d) as backbone outperforms NAS-FPN by $1.3$ AP points while using only $1/3$ FLOPs and less calculation cost.

We further measure the correlation between rewards obtained during the search process with the proxy dataset and APs attained by same architectures trained on COCO. Specifically, we randomly sample $15$ architectures from all the searched structures trained on COCO with batch size $16$. Since full training on COCO is time-consuming, we reduce the iterations to $60$K. The model is then evaluated on the COCO $2017$ validation set. As visible in Fig. 6, there is a strong correlation between search rewards and APs obtained from COCO. Poor- and well-performing architectures can be distinguished by the rewards on the proxy task very well.

3 Ablation Study

As we discussed above, it is common to use widely accepted indicators as rewards for specific tasks in the search, such as mIOU for segmentation and AP for object detection. However, we found that using AP as reward did not show a clear upward trend in short-term search rounds (blue curve in Fig. 9). We further analyze the possible reason to be that the controller tries to learn a mapping from the decoder to the reward while the calculation of AP itself is complicated, which makes it difficult to learn this mapping within a limited number of iterations. In comparison, we clearly see the increase of AP with the validation loss as RL rewards (red curve in Fig. 9).

3.2 Effectiveness of Search Space

To further discuss the impact of the search spaces $f$ and $h$, we design three experiments for verification. One is to search $f$ with the original head being fixed, one is to search $h$ with the original FPN being fixed and another is to search the entire decoder ($f$+$h$). As shown in Table 4, it turns out that searching $f$ brings slightly more benefits than searching $h$ only. And our progressive search which combines both $f$ and $h$ achieves a better result.

3.3 Impact of Deformable Convolution

As aforementioned, deformable convolutions are included in the set of candidate operations for both $f$ and $h$, which are able to adapt to the geometric variations of objects. For fair comparison, we also replace the whole standard $3\times 3$ convolutions with deformable $3\times 3$ convolutions in FPN structure of the original FCOS and repeat them twice, making the FLOPs and parameters nearly equal to our searched model. The new model is therefore called DeformFPN-FCOS. It turns out that our NAS-FCOS model still achieves better performance (AP $=38.9$ with FPN search only, and AP $=39.8$ with both FPN and Head searched) than the DeformFPN-FCOS model (AP $=38.4$) under this circumstance.

Conclusion

In this paper, we have proposed to use Neural Architecture Search to further optimize the process of designing object detection networks. It is shown in this work that top-performing detectors can be efficiently searched using carefully designed proxy tasks, search strategies and model evaluation metrics. The experiments on COCO demonstrates the efficiency of our discovered model NAS-FCOS and its flexibility to be used with various backbone architectures.

References