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

In plain words

Phones, laptops and small chips all run image models, but each has a different speed budget. The usual fix is to design and train a separate small model for every device, which costs a lot of computer time and energy when there are many devices.

This paper trains one large network once, in a way that lets smaller versions of it work well too. A smaller version is made by using fewer layers, fewer channels, smaller filters, or a smaller input image, and it reuses the big network's weights. For each device, a quick search picks the version that fits the time budget best. No new training is needed for each device.

Training is the hard part, because many versions share the same weights and can pull each other in different directions. The authors fix this by training the big version first and then gradually teaching it to work at smaller sizes. The authors report that the chosen versions match the accuracy of models trained separately, and that the total cost is much lower when many devices are involved.

TL;DR

• One shared-weight network supports sub-networks that differ in depth, expand ratio (width), kernel size and input resolution, with more than 10^19 architectures §3.2.
• Progressive shrinking trains the largest network first, then fine-tunes it on kernel, depth and width choices in turn, with distillation from the largest network §3.3 app. B.
• A specialized sub-network for a device is found by evolutionary search over an accuracy predictor and a latency lookup table, with no training per device §3.4 app. A.
• Author-reported: 80.0% ImageNet top-1 at 595M MACs under the mobile setting (<600M), after 75-epoch fine-tuning; 76.0% at 230M MACs without fine-tuning, against 75.2% for MobileNetV3-Large tab. 1 §4.2.

Problem

Deployed models have to meet different latency, energy and memory limits on different hardware. Conventional hardware-aware NAS and hand-designed models repeat the design and training for each scenario, so the total cost grows linearly with the number of scenarios §1. Pruning and slimmable networks start from one pre-designed network and shrink mainly width §2 §3.3.

The aim is a single trained model that can be specialized for many deployment scenarios at a cost that does not grow with their number §1 §3.4.

Method

Problem formulation §3.1. Let W_o be the weights of the super-network and arch_i the architectural settings of sub-network i. The training objective (the only numbered equation in the paper) is to choose W_o so that each selected sub-network C(W_o, arch_i) does about as well as a network trained alone with the same settings eq. 1. C selects the sub-network's weights from W_o.

\min_{W_o} \sum_{arch_i} \mathcal{L}_{val}\big( C(W_o, arch_i) \big) eq. 1

Architecture space §3.2. The CNN is split into units. Each unit has 2 to 4 layers (elastic depth); only the first layer of a unit has stride 2 when the map shrinks. Each layer has an expand ratio (elastic width) in {3, 4, 6} and a kernel size in {3, 5, 7} (elastic kernel). The input size runs from 128 to 224 in steps of 4 (elastic resolution). With 5 units, the paper counts roughly 2×10^19 architectures, each usable at 25 resolutions. All share one set of weights, so 7.7M parameters store all of them §3.2. The count follows from the stated sets: per unit, 3×3 choices per layer over depths 2, 3, 4 gives about 7.4×10^3 choices, and 7.4×10^3 raised to the fifth power is about 2×10^19 (our arithmetic, consistent with §3.2).

Elastic depth keeps the first D layers of a unit and drops the rest, so each depth setting corresponds to exactly one set of layers §3.2 [code: ofa/imagenet_classification/elastic_nn/networks/ofa_mbv3.py#L162-L164].

Elastic kernel size takes the centre of a 7×7 kernel as a 5×5 kernel, and the centre of that as a 3×3. Each step down applies a learned transformation matrix (25×25 from 7 to 5, 9×9 from 5 to 3), one set per layer, shared across channels. That adds 706 parameters per layer §3.3 [code: ofa/imagenet_classification/elastic_nn/modules/dynamic_op.py#L52-L65].

Elastic width sorts each layer's channels by importance, using the L1 norm of their weights. A smaller width keeps the most important channels, which share weights with the larger network §3.3 fig. 6.

Progressive shrinking §3.3. Naive training either enumerates all sub-networks each step (cost grows linearly with their number) or samples a few at random (large accuracy drops, because sub-networks interfere) §3.3. The paper's schedule trains the largest sub-network first (kernel 7, depth 4, width 6), then adds smaller choices in stages: kernel size, then depth, then width. Resolution is sampled per batch throughout §3.3 fig. 3.

Code for the schedule: [code: ofa/imagenet_classification/elastic_nn/training/progressive_shrinking.py#L287-L440] and train_ofa_net.py.

# Progressive shrinking  [§3.3] [app. B]
# Input:  trained full network (max kernel 7, depth 4, width 6)
# Output: one weight set that supports all sub-network choices

def progressive_shrinking(net, teacher, data):
    # Elastic kernel size in {3, 5, 7}, one sub-network per step  [app. B]
    # Paper: 125 epochs
    net = fine_tune(
        net, teacher, data, kernels=[3, 5, 7], subnets=1, epochs=125
    )
    # Elastic depth: first {3, 4}, then {2, 3, 4}; two sub-nets per step
    # [§3.3] [app. B]
    net = fine_tune(net, teacher, data, depths=[3, 4], subnets=2, epochs=25)
    net = fine_tune(
        net, teacher, data, depths=[2, 3, 4], subnets=2, epochs=125
    )
    # Elastic width: first {4, 6}, then {3, 4, 6}; four sub-nets per step
    # Channels are sorted by importance before this stage  [§3.3] [app. B]
    net = sort_channels(net)
    net = fine_tune(net, teacher, data, widths=[4, 6], subnets=4, epochs=25)
    net = fine_tune(
        net, teacher, data, widths=[3, 4, 6], subnets=4, epochs=125
    )
    return net

Code for one training step: [code: ofa/imagenet_classification/elastic_nn/training/progressive_shrinking.py#L157-L213] and the loss weights in train_ofa_net.py.

# One training step over several sampled sub-networks  [§3.3]
# Input:  images (B, 3, H, W); labels (B,); the largest net as teacher
# Output: one optimizer update

def train_step(net, teacher, images, labels, subnets):
    # Soft labels from the largest network  [§3.3]
    soft_labels = softmax(teacher(images))  # (B, classes)
    net.zero_grad()
    for _ in range(subnets):
        # Sample a random sub-network: kernel, width and depth per layer
        # [§3.2]
        net.sample_active_subnet()
        output = net(images)
        # Distillation term plus hard-label term; the paper gives no weights
        # [§3.3]
        loss = soft_ce(output, soft_labels) + smoothed_ce(output, labels)
        # Gradients of all sampled sub-networks are summed before one step
        # [app. B]
        loss.backward()
    optimizer.step()

Code for channel sorting: [code: ofa/imagenet_classification/elastic_nn/modules/dynamic_layers.py#L347-L413].

# Channel sorting for elastic width  [§3.3]
# Input:  expansion layer of one unit
# Output: channel order, most important first

def sort_channels(layer):
    # Importance of each middle channel: L1 norm of its outgoing weights
    # [§3.3]
    importance = sum(abs(layer.project.weight), over output channels and space)
    # Keep the widest stage's channels first, so smaller widths keep the
    # most important ones  [§3.3]
    order = argsort(importance, descending=True)
    # Reorder the expansion, depthwise and projection weights with this order
    # [§3.3]
    return order

Specialization §3.4. Once the super-network is trained, nothing is trained per device. The authors sample 16K sub-networks (architecture and input size), measure their accuracy on 10K training images, and fit an accuracy predictor to those pairs §3.4. A latency lookup table is built per target device, following ProxylessNAS §3.4. An evolutionary search then uses the two predictors to find a sub-network under the device's latency or MAC limit §3.4. The paper reports 40 GPU hours for collecting the predictor data, a cost it says stays constant as the number of deployment scenarios grows §3.4 tab. 1.

Code for the predictor and search: [code: ofa/nas/accuracy_predictor/acc_predictor.py#L15-L41] and [code: ofa/nas/search_algorithm/evolution.py#L19-L24].

# Specialize a trained super-network for one device  [§3.4] [app. A]
# Input:  super-network, device latency table, latency limit in ms
# Output: sub-network architecture and input size

def specialize(supernet, latency_table, limit_ms):
    # Sample 16K random sub-networks and measure their accuracy  [§3.4]
    samples = random_subnets(count=16000)
    accuracies = [measured_accuracy(supernet, arch) for arch in samples]
    # Fit the accuracy predictor: 3 hidden layers of 400 units  [app. A]
    predictor = fit_mlp(encode(samples), accuracies)
    # Evolutionary search with predictor and latency table; no training
    # [§3.4]
    return evolve(predictor, latency_table, limit_ms)

Hyperparameters

Items marked "code only" are not in the paper. "Not stated" means neither the paper nor the code gives the value.

Results

Author-reported. Setup: ImageNet classification §4.1. Specialized models are evaluated per device: Pixel1 phone with TF-Lite, batch size 1, for the NAS comparison §4.2 tab. 1. The baseline is MobileNetV3-Large unless stated.

¹ The text says OFA's 80.0% is 0.2 points above EfficientNet-B2, with OFA using 1.68× fewer MACs §4.2. The 79.8% figure is our derivation from that sentence, not a number the paper tabulates.

Paper vs code

Code: mit-han-lab/once-for-all at commit f03b2673db313b9167e2a1c2b7a5cad540cc1313 (MIT license, authors Han Cai and Song Han). The repo's training entry point is train_ofa_net.py; it covers kernel, depth and expand stages, not the full-network stage.

• Learning rates. The paper gives absolute LRs (0.96, 0.08, 0.24) app. B. The code sets a per-GPU base LR (3e-2, 2.5e-3, 7.5e-3) and multiplies it by the GPU count train_ofa_net.py train_ofa_net.py. The values agree only if the run used 32 GPUs. The paper states 32 GPUs for the full-network stage §4.1, not for shrinking.
• Epochs. Paper: 125 epochs for the kernel stage and for depth stage 2 app. B. Code: 120 for both train_ofa_net.py train_ofa_net.py. The width stage 2 is 125 in the paper and 120 in the code train_ofa_net.py.
• Warmup. The paper says the kernel stage uses "all other training settings … the same as training the full network" app. B. The code adds a 5-epoch warmup in the kernel stage and depth/width stage 2 train_ofa_net.py train_ofa_net.py train_ofa_net.py.
• Channel importance. The paper ranks channels by L1 norm §3.3. The code sums absolute projection weights over output channels and space [code: dynamic_layers.py#L348-L350]. It also fills group means for grouped norms dynamic_layers.py, adds stage offsets so the widest expand stage ranks first dynamic_layers.py, and reorders once at the start of the width stage progressive_shrinking.py.
• Distillation. The paper says the teacher gives soft labels from "the largest neural network" and combines two loss terms §3.3. The code uses a separate MobileNetV3Large object loaded from the D4/E6/K7 checkpoint [code: train_ofa_net.py#L207-L216] train_ofa_net.py. It weights the two terms 1.0 and 1.0 train_ofa_net.py. The paper does not give these weights, and the hard-label term uses label smoothing 0.1 [code: run_manager.py#L77-L82] that the paper does not mention.
• Validation sizes. Not stated in the paper. The code validates only at image sizes 160 and 224 and at the min/max of each choice set [code: train_ofa_net.py#L246-L253].
• Full-network stage. The paper trains it for 180 epochs at LR 2.6 §4.1. No training script for this stage is in the repo. train_ofa_net.py downloads pretrained checkpoints (D4/E6/K7 and D4/E6/K357) instead [code: train_ofa_net.py#L139-L142] train_ofa_net.py.
• Specialization fine-tuning. The 25- and 75-epoch fine-tunes of specialized sub-networks tab. 1 have no training script in the repo. The only matches for "finetune" are in evaluation and model-zoo code (eval_specialized_net.py, hubconf.py), and the released model names carry the fine-tune length [code: README.md]. I did not find the fine-tuning procedure itself.
• Predictor size. The 16K sub-networks and 10K-image measurement §3.4 are not in the repo. The accuracy predictor architecture matches app. A acc_predictor.py.
• Consistent. The kernel, depth, width and resolution sets §3.2 and the stage sample counts (1, 2, 4 sub-nets per step) match the code [code: train_ofa_net.py#L44] train_ofa_net.py train_ofa_net.py. The 706 kernel-transform parameters match the identity-initialised 25×25 and 9×9 matrices dynamic_op.py. Depth keeps the first D layers in both §3.2 and code ofa_mbv3.py.

Pitfalls

• Learning-rate scaling. The paper's LRs are only reproducible with 32 GPUs at the code's per-GPU base values. On other GPU counts, scale LR by N/32 yourself. The paper never says so.
• Batch size for shrinking. The paper gives 2048 total for the full-network stage §4.1 and does not give the batch size for shrinking. The code uses 64 per GPU train_ofa_net.py.
• Gradient aggregation. The code runs backward() once per sampled sub-network and then one optimizer step [code: progressive_shrinking.py#L170-L213]. That sums gradients without dividing by the sub-network count. The paper says "aggregate their gradients" app. B and does not say how the loss is averaged.
• Channel ordering is not plain L1. Implementing the paper's sentence "L1 norm of a channel's weight" literally misses the stage offsets and group averaging in the code dynamic_layers.py.
• Kernel transforms start from identity. Each transform matrix is initialised as the identity dynamic_op.py, so the smaller kernel starts as a plain centre crop. The paper says the matrices make up for the roles centre sub-kernels play §3.3, but it does not say how they are initialised.
• Resolution sampling. Sizes are drawn per batch and run from 128 to 224 in steps of 4 §3.3 my_random_resize_crop.py.
• Abstract vs body. The abstract claims up to 4.0% top-1 over MobileNetV3 "on diverse edge devices". I did not find that number in the body text. The body's comparable figure is 0.8 points over MobileNetV3-Large at similar Pixel1 latency §4.2 tab. 1. The per-device curves are only in fig. 10.
• FPGA numbers disagree. The text says arithmetic intensity is 48% and 43% higher than MobileNetV2 and MnasNet §4.2 fig. 12. The roofline caption says 33% and 43% against MnasNet on the two FPGAs fig. 13. Use the caption only with a stated comparison.
• Fine-tuning is part of the headline number. The 80.0%, 76.4% and 76.9% results need the 25- and 75-epoch fine-tunes tab. 1, and those procedures are not in the repo (see Paper vs code). Only 76.0% needs no fine-tune. The repo's released models are named with their fine-tune length README.md.
• Table cost units. Rows with fine-tuning list training cost as "1200 + 25N" and "1200 + 75N" GPU hours tab. 1. The paper does not say what one unit of N means in GPU hours, so these totals are approximate.
• Latency is device-specific. Latency tables are per device and per setting (TF-Lite batch 1 on phones; PyTorch batch 16 on Jetson TX2; batch 64 on the 1080Ti) §4.2. Results do not carry over to other runtimes.
• Search budget is not given. The evolutionary search settings are not in the paper. The code defaults (population 100, 500 iterations) are the only reference evolution.py.

Lineage

• Builds on: the MobileNetV3 architecture space and unit layout §4.1; predictor-guided search §3.4; a latency lookup table from ProxylessNAS §3.4; knowledge distillation §3.3. Progressive shrinking is described as a generalised form of network pruning that shrinks depth, kernel size and resolution as well as width §3.3.
• Compared against: MobileNetV2 tab. 1; NASNet-A, DARTS, MnasNet, FBNet-C, ProxylessNAS, SinglePathNAS, AutoSlim tab. 1; EfficientNet-B0 and B2 §4.2; Slimmable networks, which the paper contrasts as offering about four width settings §2.
• Superseded by: none known from the wiki.