Source-linked AI summary
Fast Neural Architecture Search of Compact Semantic Segmentation Models via Auxiliary Cells
Vladimir Nekrasov, Hao Chen, Chunhua Shen, Ian Reid
TL;DR
The paper addresses efficient neural architecture search for compact semantic-segmentation models under limited resources. It searches controller-generated decoder and auxiliary-cell designs, using progressive early termination, Polyak averaging, and knowledge distillation to speed training. The approach discovers architectures competitive with compact state-of-the-art models across semantic segmentation, pose estimation, and depth prediction tasks.
Problem
Neural architecture search for dense per-pixel tasks requires efficient training and search under limited computational budgets, unlike the better-studied classification and language-modelling settings.
Method
A reinforcement-learning controller searches compact encoder-decoder decoders and auxiliary cells, while staged training, early termination, Polyak averaging, and knowledge distillation accelerate evaluation.
Results
The discovered architectures achieve competitive performance with compact models, lower floating-point operation counts at the same output resolution, and real-time execution on GPUs and JetsonTX2.
Takeaways & Limitations
The methodology provides a way to search compact dense-prediction architectures within small computational and time budgets while retaining training-only auxiliary supervision.
Takeaways & Limitations
The encoder is initialized from a pretrained classification network, and the search is restricted to decoder architectures because full segmentation-network search from scratch is infeasible.
Abstract
from arXiv · showhide
Automated design of neural network architectures tailored for a specific task is an extremely promising, albeit inherently difficult, avenue to explore. While most results in this domain have been achieved on image classification and language modelling problems, here we concentrate on dense per-pixel tasks, in particular, semantic image segmentation using fully convolutional networks. In contrast to the aforementioned areas, the design choices of a fully convolutional network require several changes, ranging from the sort of operations that need to be used---e.g., dilated convolutions---to a solving of a more difficult optimisation problem. In this work, we are particularly interested in searching for high-performance compact segmentation architectures, able to run in real-time using limited resources. To achieve that, we intentionally over-parameterise the architecture during the training time via a set of auxiliary cells that provide an intermediate supervisory signal and can be omitted during the evaluation phase. The design of the auxiliary cell is emitted by a controller, a neural network with the fixed structure trained using reinforcement learning. More crucially, we demonstrate how to efficiently search for these architectures within limited time and computational budgets. In particular, we rely on a progressive strategy that terminates non-promising architectures from being further trained, and on Polyak averaging coupled with knowledge distillation to speed-up the convergence. Quantitatively, in 8 GPU-days our approach discovers a set of architectures performing on-par with state-of-the-art among compact models on the semantic segmentation, pose estimation and depth prediction tasks. Code will be made available here: https://github.com/drsleep/nas-segm-pytorch
1. Introduction
The paper targets compact, real-time neural architecture search for dense per-pixel tasks under limited computational budgets. It combines auxiliary-cell over-parameterisation with progressive training, Polyak averaging, and knowledge distillation to make search more efficient.
- The work addresses dense per-pixel prediction, which is more challenging for architecture search than image classification and language modelling and had lacked efficient training regimes.The methodology is presented as applicable beyond semantic segmentation to other per-pixel tasks.
- Compact architectures discovered by the method achieve competitive performance while requiring fewer resources and supporting real-time use on low-computing platforms.The stated target includes platforms such as JetsonTX2.
- The method searches fully convolutional encoder-decoder architectures whose decoder connectivity and cell operations are emitted by a controller.The same cell structure forms an auxiliary classifier for intermediate supervision.
- A progressive strategy terminates non-promising architectures early, while Polyak averaging and knowledge distillation accelerate convergence and improve performance estimation.Training is divided into stages, with early validation used to stop weak candidates.
- Auxiliary cells provide intermediate supervision during training and can be omitted during evaluation, enabling training-time over-parameterisation without increasing the inference architecture.The auxiliary-cell design is searched jointly with the decoder design.
2. Related Work
Related work spans evolutionary, Bayesian, reinforcement-learning, weight-sharing, progressive, and continuous-relaxation approaches to architecture search. Semantic segmentation architectures have largely been manually designed, while prior NAS adaptations mainly reuse classification search strategies.
- Traditional architecture search used evolutionary mutation and selection, while Bayesian optimisation estimated objective-function densities for hyper-parameter and architecture search.The related methods differ in how they explore candidate architectures.
- Reinforcement-learning NAS uses a controller to emit architectures and receives a reward after training, but historically required thousands of iterations and many GPU-days.These methods achieved state-of-the-art results primarily in image classification and natural language processing.
- Weight sharing, progressive complexity increases, and performance-prediction networks were proposed to reduce the computational cost of neural architecture search.Continuous-relaxation methods are also identified as an efficiency direction.
- Semantic segmentation architectures had mostly been manually designed around encoder-decoder or dilated-convolution paradigms.Combining these paradigms was reported to improve results across benchmarks.
- A concurrent semantic-segmentation NAS approach randomly searched for a single segmentation cell design by adapting image-classification NAS methods.The paper positions its search strategy within this emerging segmentation-NAS literature.
- Knowledge distillation trains a compact model using logits from a pretrained large network as an additional regulariser.The compact model is trained to mimic the teacher outputs.
3. Methodology
The methodology searches decoder architectures and jointly designed auxiliary cells for dense prediction. It accelerates evaluation through staged training, early termination, Polyak averaging, and knowledge distillation, while using auxiliary outputs for intermediate supervision.
- 3. Methodology: The method poses two goals: quickly obtaining reliable segmentation-performance estimates and improving training through over-parameterisation that is obsolete during inference.These goals organize the search and training methodology.
- 3.1.1 Search Space: The search fixes a pretrained encoder and explores the decoder’s selected feature maps, connectivity, and operations through a recurrent controller.The search is restricted to the decoder because full segmentation-network search from scratch is infeasible.
- 3.1.1 Search Space: The decoder cell is formed by repeatedly sampling input locations and operations, whose outputs are summed and reused as candidate inputs in later steps.The operation space includes convolutions, dilated and separable convolutions, pooling, skip connections, and a zero operation.
- 3.1.2 Training Strategy: The controller is trained with reinforcement learning using meta-train data for architecture training and meta-val performance as the reward.Sampled architectures provide the controller with scalar validation feedback.
- 3.1.2 Training Strategy: A two-stage procedure first trains only the decoder with fixed, precomputed encoder outputs, then trains the whole model end-to-end; weak candidates are stopped after stage-one validation.Continuation depends on comparing the current reward with the running reward mean.
- 3.2.2 Fast Training via Knowledge Distillation and Weights’ Averaging: Polyak parameter averaging and an L2 logit-matching loss to a pretrained teacher provide faster convergence and reliable performance estimates without significant overhead.The teacher outputs can be precomputed or acquired during training.
- 3.2.3 Intermediate Supervision via Auxiliary Cells: Auxiliary cells attached after main-cell summations provide intermediate supervision while leaving the main classifier output unchanged during training and testing.They may target ground truth, teacher predictions, or both, and the auxiliary design is tied to the decoder design.
4. Experiments
Experiments evaluate the search strategy, training enhancements, auxiliary-cell supervision, reward reliability, and transfer of discovered architectures across dense prediction tasks. The method achieves competitive compact-model performance while using limited search resources and enabling real-time inference.
- Experimental Setup: The search runs on PASCAL VOC augmented with BSD data, using validation rewards based on mean IoU, frequency-weighted IoU, and mean-pixel accuracy.Background pixels are excluded when computing the reward metrics.
- Search Results: RL increasingly samples higher-quality architectures over time and finds a larger fraction of decent candidates than random search.The authors note that more iterations might improve results further, but stop early to limit search time.
- Search Results: All combinations of Polyak averaging, auxiliary cells, and knowledge distillation outperform the baseline on both stages, while using all components gives the highest second-stage rewards.The comparison is based on training 140 randomly sampled architectures.
- Relation between search rewards and training performance: Search rewards reliably distinguish poor- from well-performing architectures when compared with longer-training mean IoU across 30 sampled architectures.The sample includes 10 architectures each from low-, medium-, and high-reward groups.
- Full Training Results: The discovered architectures achieve competitive performance against state-of-the-art compact models, with lower FLOPs at the same output resolution and real-time execution on GPUs and JetsonTX2.The evaluation covers semantic segmentation and reports additional strong results for pose estimation and depth prediction.
- Full Training Results: The highest-performing decoder uses multi-scale branches and cascaded convolutions with different kernel sizes, differing from ASPP’s reliance on dilation rates.The authors describe this design as more computationally efficient and more expressive.
Appendix A: Search Space
The search space builds decoder cells that reuse sampled encoder features through shared operations, skip connections, and repeated aggregation branches.
- Decoder connections: The decoder samples two encoder-layer indices, applies the same cell to both features, and sums the resulting outputs.Features with different spatial sizes are aligned by bilinear upsampling before the cell operation.
- Cell structure: Each cell begins with one operation, then repeatedly samples two prior indices and two operations whose outputs are summed.The resulting operation outputs and summation layers become available for later sampling.
- Cell structure: The cell search space permits operation choices at each position and uses skip connections from previous outputs.An example cell and its complete search space are illustrated in Figure 9.
Architecture description
The controller encodes each architecture as an integer sequence describing decoder connectivity and cell operations, with discovered structures visualized for three architectures.
- Architecture encoding: Architecture lists encode decoder connectivity as index pairs and cell configurations as operations applied to sampled feature indices.The connectivity structure specifies inputs from the sampling pool, while the cell configuration contains one initial operation and three branches.
- Decoder cell: The decoder cell applies the same operation structure to two specified layers, using skip connections from intermediate features.The displayed configuration distinguishes used paths from alternative unused paths.
- Discovered architectures: Figures 10–12 visualize three controller-discovered architectures, arch0, arch1, and arch2.Operation indices and abbreviations are defined in Table 4.
PASCAL VOC
The PASCAL VOC experiments use encoder–decoder training with specified optimization and augmentation settings, while the appendix supplies discovered architecture visualizations and per-class IoU results.
- Training setup: The training setup uses encoder and decoder learning rates of 1e-3 and 3e-3, respectively, with SGD for the encoder and Adam for the decoder.Training uses momentum 0.9 for the encoder and Adam parameters β1=0.9, β2=0.99, and ϵ=0.001 for the decoder.
- Architectures: Figures 10–12 provide visualizations of three discovered architectures used in the study.The supplied figure passages contain architecture encodings rather than performance values.
- Results: PASCAL VOC results include per-class intersection-over-union values across the evaluated models in Table 5.The supplied passage identifies the table’s contents but does not provide its numerical entries.
CityScapes
The CityScapes evaluation tests the discovered decoder designs with a ResNet-50 encoder and reports validation comparisons in Table 6.
- Training setup: CityScapes training uses SGD with momentum 0.9, batch size 20, and randomly scaled, mirrored, cropped, and padded 800×800 images.Logits are upsampled to the target-mask size for loss calculation.
- Results: Validation results are reported together with comparisons to several other networks in Table 6.The supplied passage does not include the numerical comparison values.
Pose estimation
Pose estimation uses cropped human instances with task-specific aspect ratios, resizing, augmentation, and heatmap generation. The setup follows the protocol described by Xiao et al.
- MPII crops use a fixed 1:1 aspect ratio, while COCO crops use 3:4.
- Both datasets resize crops so the longer side equals 256 pixels.
- MPII uses ±25% scale, ±30 degree rotation, and random-flip augmentation, while COCO uses ±30% scale and ±40-degree rotation.
- The model generates keypoint heatmaps at output stride 4 using Gaussian distributions with σ = 2.
Depth estimation
Depth estimation training uses separate encoder and decoder learning rates with SGD, scheduled decay, and a fixed 300-epoch schedule. Images are augmented and optimized with inverse Huber loss while missing-depth pixels are ignored, with qualitative outputs shown on the validation set.
- Training starts with learning rates of 1e-3 for the encoder and 7e-3 for the decoder, using SGD with momentum 0.9.
- Learning rates follow a Poly schedule, and training stops after 300 epochs.
- Batches contain 32 images across two 1080Ti GPUs, with scaling, mirroring, cropping, and padding to 500×500.
- Upsampled logits are optimized with inverse Huber loss while pixels lacking depth measurements are ignored.
- Qualitative validation-set results are visualized in Fig. 14.
Appendix C: JetsonTX2 runtime
The runtime analysis measures discovered architectures and Light-Weight RefineNet across input resolutions on JetsonTX2 and 1080Ti. Higher-FLOP models scale poorly with resolution, particularly on JetsonTX2.
- Runtime is measured for each discovered architecture and Light-Weight RefineNet while varying input resolution.
- Arch0 and Light-Weight RefineNet, which have more floating-point operations, do not scale well as input resolution increases.
- The poor scaling of higher-FLOP models is more pronounced on JetsonTX2.
- Figure 15 reports model runtimes on JetsonTX2 and 1080Ti using means and standard deviations over 100 passes.