Once-for-All: Train One Network and Specialize it for Efficient Deployment

Han Cai; Chuang Gan; Tianzhe Wang; Zhekai Zhang; Song Han

Source-linked AI summary

Once-for-All: Train One Network and Specialize it for Efficient Deployment

Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han

arXiv:1908.09791v5Aug 26, 2019cs.LGcs.CVstat.ML

TL;DR

Efficient deployment across diverse hardware and constraints is costly when each specialized network must be designed and retrained separately. OFA decouples training from architecture search, using progressive shrinking to support many subnetworks and select specialized models without retraining. Across hardware platforms, it improves accuracy-latency trade-offs while reducing deployment-design cost, including 80.0% ImageNet top1 accuracy under the mobile setting.

Problem
Designing and retraining specialized networks for many hardware platforms and efficiency constraints is computationally expensive and unscalable.
Method
OFA trains one network supporting elastic depth, width, kernel size, and resolution, then uses progressive shrinking and predictor-guided search to select specialized subnetworks.
Results
OFA consistently improves accuracy over hardware-aware NAS across diverse platforms and achieves 80.0% ImageNet top1 accuracy with 595M MACs under the mobile setting.
Takeaways & Limitations
OFA reduces specialized-network design cost from O(N) to O(1) while supporting many deployment scenarios with one trained network.

Abstract

from arXiv · show

We address the challenging problem of efficient inference across many devices and resource constraints, especially on edge devices. Conventional approaches either manually design or use neural architecture search (NAS) to find a specialized neural network and train it from scratch for each case, which is computationally prohibitive (causing $CO_2$ emission as much as 5 cars' lifetime) thus unscalable. In this work, we propose to train a once-for-all (OFA) network that supports diverse architectural settings by decoupling training and search, to reduce the cost. We can quickly get a specialized sub-network by selecting from the OFA network without additional training. To efficiently train OFA networks, we also propose a novel progressive shrinking algorithm, a generalized pruning method that reduces the model size across many more dimensions than pruning (depth, width, kernel size, and resolution). It can obtain a surprisingly large number of sub-networks ($> 10^{19}$) that can fit different hardware platforms and latency constraints while maintaining the same level of accuracy as training independently. On diverse edge devices, OFA consistently outperforms state-of-the-art (SOTA) NAS methods (up to 4.0% ImageNet top1 accuracy improvement over MobileNetV3, or same accuracy but 1.5x faster than MobileNetV3, 2.6x faster than EfficientNet w.r.t measured latency) while reducing many orders of magnitude GPU hours and $CO_2$ emission. In particular, OFA achieves a new SOTA 80.0% ImageNet top-1 accuracy under the mobile setting ($<$600M MACs). OFA is the winning solution for the 3rd Low Power Computer Vision Challenge (LPCVC), DSP classification track and the 4th LPCVC, both classification track and detection track. Code and 50 pre-trained models (for many devices & many latency constraints) are released at https://github.com/mit-han-lab/once-for-all.

1 INTRODUCTION

OFA addresses the cost of deploying specialized networks across diverse hardware by training one flexible network and selecting subnetworks without retraining. It reports broad hardware-aware gains, including 80.0% ImageNet top-1 accuracy under the mobile constraint.

Diverse hardware resources and efficiency constraints require different neural network architectures, making broad deployment difficult.
Designing and retraining a specialized network for every deployment scenario is expensive, with total cost growing linearly in the number of scenarios.The repeated process also increases energy consumption and CO2 emission.
OFA decouples model training from architecture search, selecting specialized subnetworks across depth, width, kernel size, and resolution without retraining.The reported design cost falls from O(N) to O(1).
Training OFA is challenging because more than 10^19 weight-sharing subnetworks must retain accuracy while interfering with one another.Progressive shrinking is introduced to address this training challenge.
80.0% ImageNet top1 accuracy is achieved under the mobile setting with fewer than 600M MACs.The comparison is presented against state-of-the-art CNN models.

2 RELATED WORK

Prior work improves efficiency through compact architectures, compression, NAS, and dynamic computation, but these approaches do not provide OFA's broad jointly supported architectural space. OFA instead uses progressive shrinking to cover diverse deployment scenarios.

Efficient Deep Learning: Efficient deep learning research includes compact CNN architectures and compression methods such as pruning and quantization.
Neural Architecture Search: NAS automates architecture design, while hardware-aware NAS incorporates inference efficiency into the search process.
Dynamic Neural Networks: Dynamic neural networks adapt computation to input images by skipping layers, exiting early, pruning channels, or reducing feature-map resolution.
OFA derives specialized networks for many deployment scenarios rather than optimizing on top of one existing architecture.This expands the available optimization space but makes training more challenging.

3 METHOD

OFA formalizes a weight-sharing network that selects subnetworks across multiple architectural dimensions, then trains it with progressive shrinking and specializes it through predictor-guided search. The method supports a large subnetwork space while keeping specialization training-free.

Problem Formalization: The selection scheme C(W_o, arch_i) extracts a subnetwork with configuration arch_i from the once-for-all weights W_o.The objective is to preserve independently trained accuracy across supported subnetworks.
Once-for-All Network: OFA supports elastic depth, width, kernel size, and input resolution within one weight-sharing model.The architecture space varies layer counts, channel expansion, kernel sizes, and image sizes.
Naive Training: Naive enumeration is too costly, while sampling a few subnetworks causes accuracy drops because shared subnetworks interfere during training.
Progressive Shrinking: Progressive shrinking trains the largest network first, then progressively fine-tunes smaller nested subnetworks using shared important weights.This ordering reduces interference and improves initialization for smaller models.
Elastic Dimensions: Elastic depth keeps the first D layers, elastic width selects channels with larger L1 norms, and elastic kernels use nested centered subkernels.These mechanisms share weights between larger and smaller subnetworks.
Specialized Deployment: After OFA training, accuracy predictors and hardware-specific latency lookup tables guide subnetwork search without additional training.The accuracy predictor is trained from 16K sampled subnetworks and their validation accuracies.
Specialized Deployment: 40 GPU hours collect predictor data, with this cost remaining constant regardless of the number of deployment scenarios.

4 EXPERIMENTS

Experiments evaluate OFA across ImageNet, diverse hardware platforms, and latency constraints, showing strong accuracy–efficiency trade-offs and substantially lower deployment-design cost than hardware-aware NAS.

Experimental setup: OFA is evaluated on ImageNet across mobile, cloud, embedded, CPU, GPU, and FPGA platforms with different latency constraints.The study includes Samsung, Google, and LG phones, NVIDIA GPUs, Jetson TX2, Intel Xeon, and Xilinx FPGAs.
Training cost: 1,200 GPU hours is the one-time OFA training cost, which can be amortized across many deployment scenarios.The reported training used 32 GPUs and the full network was subsequently fine-tuned.
Training the once-for-all network: 74.8% top1 accuracy at 226M MACs with progressive shrinking matches MobileNetV3-Large, versus 71.5% without progressive shrinking.The comparison uses architecture setting (D=4, W=3, K=3).
Mobile devices: With N = 40 deployment scenarios, OFA emits 16× less CO2 than ProxylessNAS, 19× less than FBNet, and 1,300× less than MnasNet.Without retraining, OFA reaches 76.0% top1 accuracy on ImageNet, 0.8% higher than MobileNetV3-Large at similar mobile latency.
Resource constraints: 80.0% ImageNet top1 accuracy at 595M MACs is achieved under the mobile setting, while 80.1% at 143ms Pixel1 latency is 2.6× faster than EfficientNet-B2.At 389M MACs, OFA reaches 79.1% top1 accuracy, 2.8% above EfficientNet-B0 with similar MACs.
Diverse hardware platforms: With similar latency as MobileNetV2 0.35, OFA #25 raises ImageNet top1 accuracy from 60.3% to 72.6% on the 1080Ti GPU.The improvement is +12.3% on that GPU comparison.
Specialized hardware accelerators: OFA increases arithmetic intensity by 48%/43% over MobileNetV2 and MnasNet and raises GOPS/s by 70%-90% on the evaluated FPGAs.The FPGA analysis uses Xilinx Vitis AI with 8-bit quantization.

5 CONCLUSION

The conclusion presents OFA as a unified deployment methodology that separates training from search, supports elastic architectures, and reduces the cost of specialization.

Conclusion: OFA decouples model training from architecture search for efficient deployment across many hardware platforms.The methodology selects specialized subnetworks from a trained once-for-all network.
Conclusion: Elastic depth, width, kernel size, and resolution let one once-for-all network support different architectural configurations.These configurations replace designing and training a separate network for each deployment scenario.
Conclusion: Progressive shrinking reduces training cost by enabling many subnetworks to reach the same accuracy level as independently trained models.The method addresses interference among subnetworks of different sizes during joint optimization.
Conclusion: Experiments across diverse hardware platforms and efficiency constraints demonstrate OFA’s effectiveness as an automated design methodology with hardware cost models in the loop.The conclusion frames OFA as an ecosystem for efficient neural-network design.

A DETAILS OF THE ACCURACY PREDICTOR

The accuracy predictor encodes a candidate architecture and input image size as a vector, then predicts the selected subnetworks’ accuracy for search.

Architecture encoding: Each layer is encoded with a one-hot vector for kernel size and expansion ratio, while skipped layers receive zero vectors.An additional one-hot vector represents input image size.
Accuracy prediction: The concatenated architecture and image-size representation is used by a three-layer feedforward neural network with 400 hidden units per layer.The predictor supports accuracy evaluation for selected subnetworks.
Accuracy prediction: The plotted quantity is the accuracy of selected subnetworks.The figure labels this measure as “Acc of Selected Sub-net (%)”.

B IMPLEMENTATION DETAILS OF PROGRESSIVE SHRINKING

Progressive shrinking is implemented through sequential fine-tuning stages that add elastic kernel size, depth, and width by sampling subnetworks and aggregating gradients.

Kernel size: Elastic kernel size is added first by sampling one subnetwork per update for 125 fine-tuning epochs.The kernel-size choices are K ∈ [7, 5, 3].
Depth: Elastic depth is added in two stages using two sampled subnetworks and aggregated gradients per update.The stages use D ∈ [4, 3] and then D ∈ [4, 3, 2], lasting 25 and 125 epochs respectively.
Width: Elastic width is added in two stages using four sampled subnetworks and aggregated gradients per update.The stages use W ∈ [6, 4] and then W ∈ [6, 4, 3], lasting 25 and 125 epochs respectively.