Auto-DeepLab: Hierarchical Neural Architecture Search for Semantic Image Segmentation

Introduction

Deep neural networks have been proved successful across a large variety of artificial intelligence tasks, including image recognition , speech recognition , machine translation etc. While better optimizers and better normalization techniques certainly played an important role, a lot of the progress comes from the design of neural network architectures. In computer vision, this holds true for both image classification and dense image prediction .

More recently, in the spirit of AutoML and democratizing AI, there has been significant interest in designing neural network architectures automatically, instead of relying heavily on expert experience and knowledge. Importantly, in the past year, Neural Architecture Search (NAS) has successfully identified architectures that exceed human-designed architectures on large-scale image classification problems .

Image classification is a good starting point for NAS, because it is the most fundamental and well-studied high-level recognition task. In addition, there exists benchmark datasets (e.g., CIFAR-10) with relatively small images, resulting in less computation and faster training. However, image classification should not be the end point for NAS, and the current success shows promise to extend into more demanding domains. In this paper, we study Neural Architecture Search for semantic image segmentation, an important computer vision task that assigns a label like “person” or “bicycle” to each pixel in the input image.

Naively porting ideas from image classification would not suffice for semantic segmentation. In image classification, NAS typically applies transfer learning from low resolution images to high resolution images , whereas optimal architectures for semantic segmentation must inherently operate on high resolution imagery. This suggests the need for: (1) a more relaxed and general search space to capture the architectural variations brought by the higher resolution, and (2) a more efficient architecture search technique as higher resolution requires heavier computation.

We notice that modern CNN designs usually follow a two-level hierarchy, where the outer network level controls the spatial resolution changes, and the inner cell level governs the specific layer-wise computations. The vast majority of current works on NAS follow this two-level hierarchical design, but only automatically search the inner cell level while hand-designing the outer network level. This limited search space becomes problematic for dense image prediction, which is sensitive to the spatial resolution changes. Therefore in our work, we propose a trellis-like network level search space that augments the commonly-used cell level search space first proposed in to form a hierarchical architecture search space. Our goal is to jointly learn a good combination of repeatable cell structure and network structure specifically for semantic image segmentation.

In terms of the architecture search method, reinforcement learning and evolutionary algorithms tend to be computationally intensive even on the low resolution CIFAR-10 dataset, therefore probably not suitable for semantic image segmentation. We draw inspiration from the differentiable formulation of NAS , and develop a continuous relaxation of the discrete architectures that exactly matches the hierarchical architecture search space. The hierarchical architecture search is conducted via stochastic gradient descent. When the search terminates, the best cell architecture is decoded greedily, and the best network architecture is decoded efficiently using the Viterbi algorithm. We directly search architecture on $321\times 321$ image crops from Cityscapes . The search is very efficient and only takes about $3$ days on one P100 GPU.

We report experimental results on multiple semantic segmentation benchmarks, including Cityscapes , PASCAL VOC 2012 , and ADE20K . Without ImageNet pretraining, our best model significantly outperforms FRRN-B by $8.6\%$ and GridNet by $10.9\%$ on Cityscapes test set, and performs comparably with other ImageNet-pretrained state-of-the-art models when also exploiting the coarse annotations on Cityscapes. Notably, our best model (without pretraining) attains the same performance as DeepLabv3+ (with pretraining) while being $2.23$ times faster in Multi-Adds. Additionally, our light-weight model attains the performance only $1.2\%$ lower than DeepLabv3+ , while requiring $76.7\%$ fewer parameters and being $4.65$ times faster in Multi-Adds. Finally, on PASCAL VOC 2012 and ADE20K, our best model outperforms several state-of-the-art models while using strictly less data for pretraining.

To summarize, the contribution of our paper is four-fold:

Ours is one of the first attempts to extend NAS beyond image classification to dense image prediction.

We propose a network level architecture search space that augments and complements the much-studied cell level one, and consider the more challenging joint search of network level and cell level architectures.

We develop a differentiable, continuous formulation that conducts the two-level hierarchical architecture search efficiently in $3$ GPU days.

Without ImageNet pretraining, our model significantly outperforms FRRN-B and GridNet, and attains comparable performance with other ImageNet-pretrained state-of-the-art models on Cityscapes. On PASCAL VOC 2012 and ADE20K, our best model also outperforms several state-of-the-art models.

Related Work

Convolutional neural networks deployed in a fully convolutional manner (FCNs ) have achieved remarkable performance on several semantic segmentation benchmarks. Within the state-of-the-art systems, there are two essential components: multi-scale context module and neural network design. It has been known that context information is crucial for pixel labeling tasks . Therefore, PSPNet performs spatial pyramid pooling at several grid scales (including image-level pooling ), while DeepLab applies several parallel atrous convolution with different rates. On the other hand, the improvement of neural network design has significantly driven the performance from AlexNet , VGG , Inception , ResNet to more recent architectures, such as Wide ResNet , ResNeXt , DenseNet and Xception . In addition to adopting those networks as backbones for semantic segmentation, one could employ the encoder-decoder structures which efficiently captures the long-range context information while keeping the detailed object boundaries. Nevertheless, most of the models require initialization from the ImageNet pretrained checkpoints except FRRN and GridNet for the task of semantic segmentation. Specifically, FRRN employs a two-stream system, where full-resolution information is carried in one stream and context information in the other pooling stream. GridNet, building on top of a similar idea, contains multiple streams with different resolutions. In this work, we apply neural architecture search for network backbones specific for semantic segmentation. We further show state-of-the-art performance without ImageNet pretraining, and significantly outperforms FRRN and GridNet on Cityscapes .

Neural Architecture Search aims at automatically designing neural network architectures, hence minimizing human hours and efforts. While some works search RNN cells for language tasks, more works search good CNN architectures for image classification.

Several papers used reinforcement learning (either policy gradients or Q-learning ) to train a recurrent neural network that represents a policy to generate a sequence of symbols specifying the CNN architecture. An alternative to RL is to use evolutionary algorithms (EA), that “evolves” architectures by mutating the best architectures found so far . However, these RL and EA methods tend to require massive computation during the search, usually thousands of GPU days. PNAS proposed a progressive search strategy that markedly reduced the search cost while maintaining the quality of the searched architecture. NAO embedded architectures into a latent space and performed optimization before decoding. Additionally, several works utilized architectural sharing among sampled models instead of training each of them individually, thereby further reduced the search cost. Our work follows the differentiable NAS formulation and extends it into the more general hierarchical setting.

Earlier papers, e.g., , tried to directly construct the entire network. However, more recent papers have shifted to searching the repeatable cell structure, while keeping the outer network level structure fixed by hand. First proposed in , this strategy is likely inspired by the two-level hierarchy commonly used in modern CNNs.

Our work still uses this cell level search space to keep consistent with previous works. Yet one of our contributions is to propose a new, general-purpose network level search space, since we wish to jointly search across this two-level hierarchy. Our network level search space shares a similar outlook as , but the important difference is that kept the entire “fabrics” with no intention to alter the architecture, whereas we associate an explicit weight for each connection and focus on decoding a single discrete structure. In addition, was evaluated on segmenting face images into $3$ classes , whereas our models are evaluated on large-scale segmentation datasets such as Cityscapes , PASCAL VOC 2012 , and ADE20K .

The most similar work to ours is , which also studied NAS for semantic image segmentation. However, focused on searching the much smaller Atrous Spatial Pyramid Pooling (ASPP) module using random search, whereas we focus on searching the much more fundamental network backbone architecture using more advanced and more efficient search methods.

Architecture Search Space

This section describes our two-level hierarchical architecture search space. For the inner cell level (Sec. 3.1), we reuse the one adopted in to keep consistent with previous works. For the outer network level (Sec. 3.2), we propose a novel search space based on observation and summarization of many popular designs.

We define a cell to be a small fully convolutional module, typically repeated multiple times to form the entire neural network. More specifically, a cell is a directed acyclic graph consisting of $B$ blocks.

Each block is a two-branch structure, mapping from $2$ input tensors to $1$ output tensor. Block $i$ in cell $l$ may be specified using a $5$-tuple $(I_{1},I_{2},O_{1},O_{2},C)$, where $I_{1},I_{2}\in\mathcal{I}_{i}^{l}$ are selections of input tensors, $O_{1},O_{2}\in\mathcal{O}$ are selections of layer types applied to the corresponding input tensor, and $C\in\mathcal{C}$ is the method used to combine the individual outputs of the two branches to form this block’s output tensor, $H_{i}^{l}$. The cell’s output tensor $H^{l}$ is simply the concatenation of the blocks’ output tensors $H_{1}^{l},\ldots,H_{B}^{l}$ in this order.

The set of possible input tensors, $\mathcal{I}_{i}^{l}$, consists of the output of the previous cell $H^{l-1}$, the output of the previous-previous cell $H^{l-2}$, and previous blocks’ output in the current cell $\{H_{1}^{l},\ldots,H_{i}^{l}\}$. Therefore, as we add more blocks in the cell, the next block has more choices as potential source of input.

The set of possible layer types, $\mathcal{O}$, consists of the following $8$ operators, all prevalent in modern CNNs:

For the set of possible combination operators $\mathcal{C}$, we simply let element-wise addition to be the only choice.

2 Network Level Search Space

In the image classification NAS framework pioneered by , once a cell structure is found, the entire network is constructed using a pre-defined pattern. Therefore the network level was not part of the architecture search, hence its search space has never been proposed nor designed.

This pre-defined pattern is simple and straightforward: a number of “normal cells” (cells that keep the spatial resolution of the feature tensor) are separated equally by inserting “reduction cells” (cells that divide the spatial resolution by $2$ and multiply the number of filters by $2$). This keep-downsampling strategy is reasonable in the image classification case, but in dense image prediction it is also important to keep high spatial resolution, and as a result there are more network level variations .

Among the various network architectures for dense image prediction, we notice two principles that are consistent:

The spatial resolution of the next layer is either twice as large, or twice as small, or remains the same.

The smallest spatial resolution is downsampled by $32$.

Following these common practices, we propose the following network level search space. The beginning of the network is a two-layer “stem” structure that each reduces the spatial resolution by a factor of $2$. After that, there are a total of $L$ layers with unknown spatial resolutions, with the maximum being downsampled by $4$ and the minimum being downsampled by $32$. Since each layer may differ in spatial resolution by at most $2$, the first layer after the stem could only be either downsampled by $4$ or $8$. We illustrate our network level search space in Fig. 1. Our goal is then to find a good path in this $L$-layer trellis.

In Fig. 2 we show that our search space is general enough to cover many popular designs. In the future, we have plans to relax this search space even further to include U-net architectures , where layer $l$ may receive input from one more layer preceding $l$ in addition to $l-1$.

We reiterate that our work searches the network level architecture in addition to the cell level architecture. Therefore our search space is strictly more challenging and general-purpose than previous works.

Methods

We begin by introducing a continuous relaxation of the (exponentially many) discrete architectures that exactly matches the hierarchical architecture search described above. We then discuss how to perform architecture search via optimization, and how to decode back a discrete architecture after the search terminates.

We reuse the continuous relaxation described in . Every block’s output tensor $H_{i}^{l}$ is connected to all hidden states in $\mathcal{I}_{i}^{l}$:

In addition, we approximate each $O_{j\rightarrow i}$ with its continuous relaxation $\bar{O}_{j\rightarrow i}$, defined as:

In other words, $\alpha_{j\rightarrow i}^{k}$ are normalized scalars associated with each operator $O^{k}\in\mathcal{O}$, easily implemented as softmax.

Recall from Sec. 3.1 that $H^{l-1}$ and $H^{l-2}$ are always included in $\mathcal{I}_{i}^{l}$, and that $H^{l}$ is the concatenation of $H_{1}^{l},\ldots,H_{B}^{l}$. Together with Eq. (1) and Eq. (2), the cell level update may be summarized as:

1.2 Network Architecture

We design the network level continuous relaxation to exactly match the search space described in Sec. 3.2. We associated a scalar with each gray arrow in Fig. 1, and the network level update is:

where $s=4,8,16,32$ and $l=1,2,\ldots,L$. The scalars $\beta$ are normalized such that

Eq. (6) shows how the continuous relaxations of the two-level hierarchy are weaved together. In particular, $\beta$ controls the outer network level, hence depends on the spatial size and layer index. Each scalar in $\beta$ governs an entire set of $\alpha$, yet $\alpha$ specifies the same architecture that depends on neither spatial size nor layer index.

As illustrated in Fig. 1, Atrous Spatial Pyramid Pooling (ASPP) modules are attached to each spatial resolution at the $L$-th layer (atrous rates are adjusted accordingly). Their outputs are bilinear upsampled to the original resolution before summed to produce the prediction.

2 Optimization

The advantage of introducing this continuous relaxation is that the scalars controlling the connection strength between different hidden states are now part of the differentiable computation graph. Therefore they can be optimized efficiently using gradient descent. We adopt the first-order approximation in , and partition the training data into two disjoint sets trainA and trainB. The optimization alternates between:

Update network weights $w$ by $\nabla_{w}\mathcal{L}_{trainA}(w,\alpha,\beta)$

Update architecture $\alpha,\beta$ by $\nabla_{\alpha,\beta}\mathcal{L}_{trainB}(w,\alpha,\beta)$

where the loss function $\mathcal{L}$ is the cross entropy calculated on the semantic segmentation mini-batch. The disjoint set partition is to prevent the architecture from overfitting the training data.

3 Decoding Discrete Architectures

Following , we decode the discrete cell architecture by first retaining the 2 strongest predecessors for each block (with the strength from hidden state $j$ to hidden state $i$ being $\max_{k,O^{k}\neq zero}\alpha_{j\rightarrow i}^{k}$; recall from Sec. 3.1 that “zero” means “no connection”), and then choose the most likely operator by taking the argmax.

Eq. (7) essentially states that the “outgoing probability” at each of the blue nodes in Fig. 1 sums to $1$. In fact, the $\beta$ values can be interpreted as the “transition probability” between different “states” (spatial resolution) across different “time steps” (layer number). Quite intuitively, our goal is to find the path with the “maximum probability” from start to end. This path can be decoded efficiently using the classic Viterbi algorithm, as in our implementation.

Experimental Results

Herein, we report our architecture search implementation details as well as the search results. We then report semantic segmentation results on benchmark datasets with our best found architecture.

We consider a total of $L=12$ layers in the network, and $B=5$ blocks in a cell. The network level search space has $2.9\times 10^{4}$ unique paths, and the number of cell structures is $5.6\times 10^{14}$. So the size of the joint, hierarchical search space is in the order of $10^{19}$.

We follow the common practice of doubling the number of filters when halving the height and width of feature tensor. Every blue node in Fig. 1 with downsample rate $s$ has $B\times F\times\frac{s}{4}$ output filters, where $F$ is the filter multiplier controlling the model capacity. We set $F=8$ during the architecture search. A stride $2$ convolution is used for all $\frac{s}{2}\rightarrow s$ connections, both to reduce spatial size and double the number of filters. Bilinear upsampling followed by $1\times 1$ convolution is used for all $2s\rightarrow s$ connections, both to increase spatial size and halve the number of filters.

The Atrous Spatial Pyramid Pooling module used in has $5$ branches: one $1\times 1$ convolution, three $3\times 3$ convolution with various atrous rates, and pooled image feature. During the search, we simplify ASPP to have $3$ branches instead of $5$ by only using one $3\times 3$ convolution with atrous rate $\frac{96}{s}$. The number of filters produced by each ASPP branch is still $B\times F\times\frac{s}{4}$.

We conduct architecture search on the Cityscapes dataset for semantic image segmentation. More specifically, we use $321\times 321$ random image crops from half-resolution ($512\times 1024$) images in the train_fine set. We randomly select half of the images in train_fine as trainA, and the other half as trainB (see Sec. 4.2).

The architecture search optimization is conducted for a total of $40$ epochs. The batch size is $2$ due to GPU memory constraint. When learning network weights $w$, we use SGD optimizer with momentum $0.9$, cosine learning rate that decays from $0.025$ to $0.001$, and weight decay $0.0003$. The initial values of $\alpha,\beta$ before softmax are sampled from a standard Gaussian times $0.001$. They are optimized using Adam optimizer with learning rate $0.003$ and weight decay $0.001$. We empirically found that if $\alpha,\beta$ are optimized from the beginning when $w$ are not well trained, the architecture tends to fall into bad local optima. Therefore we start optimizing $\alpha,\beta$ after $20$ epochs. The entire architecture search optimization takes about $3$ days on one P100 GPU. Fig. 4 shows that the validation accuracy steadily improves throughout this process. We also tried searching for longer epochs ($60$, $80$, $100$), but did not observe benefit.

Fig. 3 visualizes the best architecture found. In terms of network level architecture, higher resolution is preferred at both beginning (stays at downsample by $4$ for longer) and end (ends at downsample by $8$). We also show the strongest outgoing connection at each node using gray dashed arrows. We observe a general tendency to downsample in the first $3/4$ layers and upsample in the last $1/4$ layers. In terms of cell level architecture, the conjunction of atrous convolution and depthwise-separable convolution is often used, suggesting that the importance of context has been learned. Note that atrous convolution is rarely found to be useful in cells for image classificationAmong NASNet-{A, B, C}, PNASNet-{1, 2, 3, 4, 5}, AmoebaNet-{A, B, C}, ENAS, DARTS, atrous convolution was used only once in AmoebaNet-B reduction cell..

2 Semantic Segmentation Results

We evaluate the performance of our found best architecture (Fig. 3) on Cityscapes , PASCAL VOC 2012 , and ADE20K datasets.

We follow the same training protocol in . In brief, during training we adopt a polynomial learning rate schedule with initial learning rate $0.05$, and large crop size (e.g., $769\times 769$ on Cityscapes, and $513\times 513$ on PASCAL VOC 2012 and resized ADE20K images). Batch normalization parameters are fine-tuned during training. The models are trained from scratch with 1.5M iterations on Cityscapes, 1.5M iterations on PASCAL VOC 2012, and 4M iterations on ADE20K, respectively.

We adopt the simple encoder-decoder structure similar to DeepLabv3+ . Specifically, our encoder consists of our found best network architecture augmented with the ASPP module , and our decoder is the same as the one in DeepLabv3+ which recovers the boundary information by exploiting the low-level features that have downsample rate $4$. Additionally, we redesign the “stem” structure with three $3\times 3$ convolutions (with stride $2$ in the first and third convolutions). The first two convolutions have $64$ filters while the third convolution has $128$ filters. This “stem” has been shown to be effective for segmentation in .

Cityscapes contains high quality pixel-level annotations of $5000$ images with size $1024\times 2048$ ($2975$, $500$, and $1525$ for the training, validation, and test sets respectively) and about $20000$ coarsely annotated training images. Following the evaluation protocol , $19$ semantic labels are used for evaluation without considering the void label.

In Tab. 2, we report the Cityscapes validation set results. Similar to MobileNets , we adjust the model capacity by changing the filter multiplier $F$. As shown in the table, higher model capacity leads to better performance at the cost of slower speed (indicated by larger Multi-Adds).

In Tab. 3, we show that increasing the training iterations from 500K to 1.5M iterations improves the performance by $2.8\%$, when employing our light-weight model variant, Auto-DeepLab-S. Additionally, adopting the Scheduled Drop Path further improves the performance by $1.74\%$, reaching $79.74\%$ on Cityscapes validation set.

We then report the test set results in Tab. 4. Without any pretraining, our best model (Auto-DeepLab-L) significantly outperforms FRNN-B by $8.6\%$ and GridNet by $10.9\%$. With extra coarse annotations, our model Auto-DeepLab-L, without pretraining on ImageNet , achieves the test set performance of $82.1\%$, outperforming PSPNet and Mapillary , and attains the same performance as DeepLabv3+ while requiring $55.2\%$ fewer Mutli-Adds computations. Notably, our light-weight model variant, Auto-DeepLab-S, attains $80.9\%$ on the test set, comparable to PSPNet, while using merely 10.15M parameters and 333.25B Multi-Adds.

2.2 PASCAL VOC 2012

PASCAL VOC 2012 contains $20$ foreground object classes and one background class. We augment the original dataset with the extra annotations provided by , resulting in $10582$ (train_aug) training images.

In Tab. 5, we report our validation set results. Our best model, Auto-DeepLab-L, with single scale inference significantly outperforms by $20.36\%$. Additionally, for all our model variants, adopting multi-scale inference improves the performance by about $1\%$. Further pretraining our models on COCO for 4M iterations improves the performance significantly.

Finally, we report the PASCAL VOC 2012 test set result with our COCO-pretrained model variants in Tab. 6. As shown in the table, our best model attains the performance of $85.6\%$ on the test set, outperforming RefineNet and PSPNet . Our model is lagged behind the top-performing DeepLabv3+ with Xception-65 as network backbone by $2.2\%$. We think that PASCAL VOC 2012 dataset is too small to train models from scratch and pretraining on ImageNet is still beneficial in this case.

2.3 ADE20K

ADE20K has $150$ semantic classes and high quality annotations of $20000$ training images and $2000$ validation images. In our experiments, the images are all resized so that the longer side is $513$ during training.

In Tab. 7, we report our validation set results. Our models outperform some state-of-the-art models, including RefineNet , UPerNet , and PSPNet (ResNet-152) ; however, without any ImageNet pretraining, our performance is lagged behind the latest work of .

Conclusion

In this paper, we present one of the first attempts to extend Neural Architecture Search beyond image classification to dense image prediction problems. Instead of fixating on the cell level, we acknowledge the importance of spatial resolution changes, and embrace the architectural variations by incorporating the network level into the search space. We also develop a differentiable formulation that allows efficient (about $1000\times$ faster than DPC ) architecture search over our two-level hierarchical search space. The result of the search, Auto-DeepLab, is evaluated by training on benchmark semantic segmentation datasets from scratch. On Cityscapes, Auto-DeepLab significantly outperforms the previous state-of-the-art by $8.6\%$, and performs comparably with ImageNet-pretrained top models when exploiting the coarse annotations. On PASCAL VOC 2012 and ADE20K, Auto-DeepLab also outperforms several ImageNet-pretrained state-of-the-art models.

For future work, within the current framework, related applications such as object detection should be plausible; we could also try untying the cell architecture $\alpha$ across different layers (cf. ) with little computation overhead. Beyond the current framework, a more general network level search space should be beneficial (cf. Sec. 3.2).

We thank Sergey Ioffe for valuable feedback; Cloud AI and Mobile Vision team for support. CL and AY acknowledge a gift from YiTu.

References