[Paper Review] Once-for-All: Train One Network and Specialize it for Efficient Deployment

Han Cai, Chuang Gan+3 more|arXiv (Cornell University)|Aug 26, 2019

Advanced Neural Network Applications38 references679 citations

TL;DR

OFA trains a single flexible network that can specialize into many sub-networks for diverse hardware with no retraining, enabling efficient deployment across devices while maintaining accuracy.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for efficient deployment across diverse hardware with minimal retraining and cost.
• Introduce a single once-for-all network that supports many architectural configurations (depth, width, kernel size, resolution).
• Propose a training regime that yields accurate sub-networks without retraining for each deployment scenario.

Proposed method

• Define an elastic architecture space with depth, width, kernel size, and resolution to map to sub-networks.
• Train the large OFA network with progressive shrinking to gradually support smaller sub-networks while sharing weights.
• Use knowledge distillation to stabilize training across nested sub-networks.
• During specialization, build neural-network-twins (accuracy predictor and latency lookup) to guide an evolutionary search for the best sub-network per hardware constraint.
• Decouple training from search to reduce cost from O(N) to O(1) across deployment scenarios.

Experimental results

Research questions

• RQ1 Can a single OFA network support a huge number of sub-networks ( >10^19 ) while preserving accuracy comparable to independently trained nets?
• RQ2 Does progressive shrinking effectively mitigate interference among sub-networks during joint training?
• RQ3 Can a predictor-guided search (neural-network twins) efficiently identify optimal sub-networks for diverse hardware with negligible cost?
• RQ4 How does OFA perform versus state-of-the-art hardware-aware NAS methods across cloud and edge devices in terms of accuracy, latency, and energy footprint?

Key findings

• OFA achieves superior accuracy-latency trade-offs across multiple hardware platforms compared with SOTA hardware-aware NAS methods.
• On ImageNet mobile setting (<600M MACs), OFA reaches 80.0% top-1 accuracy with 595M MACs, a new mobile SOTA.
• OFA reduces training and design costs by orders of magnitude and lowers CO2 emissions compared with NAS approaches when supporting many deployment scenarios.
• Progressive shrinking enables a large sub-network space (>10^19 architectures) to be trained efficiently, maintaining accuracy comparable to independently trained sub-networks.
• Specialized OFA sub-networks at various devices (CPU, GPU, FPGA, mobile) consistently surpass MobileNetV2/MnasNet/others at similar latency, while requiring negligible additional training for new hardware.