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

TL;DR

Proposes Once-for-All (OFA) network supporting diverse architectures, achieving 80.0% ImageNet top-1 accuracy with low complexity.

cs.LG 🔴 Advanced 2019-08-27 44 views
Han Cai Chuang Gan Tianzhe Wang Zhekai Zhang Song Han
Neural Architecture Search Model Compression Edge Deployment Multi-dimensional Pruning Efficient Deep Learning

Key Findings

Methodology

The paper introduces a progressive shrinking algorithm that supports multi-dimensional pruning—depth, width, kernel size, and input resolution—to generate over 10^19 sub-networks. Training involves shared weights and knowledge distillation, optimizing the entire architecture space jointly. A predictor-guided search efficiently matches sub-networks to hardware constraints, decoupling training from deployment. This approach reduces computational costs by orders of magnitude while maintaining high accuracy. The core components include the multi-dimensional elastic architecture space, progressive shrinking training scheme, and predictor-based architecture search, enabling rapid customization for various hardware platforms.

Key Results

  • On ImageNet, OFA achieves 80.0% top-1 accuracy under mobile constraints (<600M MACs), outperforming MobileNetV3 by 4.0%, with 1.5× faster inference. GPU training time is reduced by over 50×, with significant CO2 emission savings.
  • Across diverse hardware (GPU, FPGA, mobile phones), OFA supports multiple architectures with latency and accuracy trade-offs. Experimental results show speedups of 1.5-2.6× over independently trained models, with comparable or better accuracy.
  • The progressive shrinking method effectively maintains accuracy across a vast architecture space, validating its ability to adapt to different deployment scenarios while drastically reducing training costs.

Significance

This work addresses the scalability bottleneck of traditional NAS and model compression by enabling a single trained network to serve multiple deployment scenarios. It significantly reduces the time, energy, and resource costs associated with model development, making high-performance deep learning models more accessible for edge devices and large-scale industrial applications. The decoupling of training and search processes fosters rapid deployment and customization, advancing the practical deployment of AI in resource-constrained environments.

Technical Contribution

The paper introduces a multi-dimensional progressive shrinking algorithm that supports elastic depth, width, kernel size, and resolution. It leverages shared weights and knowledge distillation to train a single, versatile network capable of supporting over 10^19 sub-architectures. Predictor-guided architecture search accelerates deployment customization, while the decoupling of training and search reduces computational overhead. These innovations collectively enable scalable, hardware-aware neural network design, pushing the boundaries of efficient deep learning.

Novelty

This is the first work to unify multi-dimensional elastic architecture design with progressive shrinking, supporting an unprecedentedly large architecture space (>10^19 configurations). Unlike prior methods limited to single-dimension adjustments or small model families, this approach offers a comprehensive, scalable solution for multi-scenario deployment, fundamentally changing how neural networks are trained and adapted.

Limitations

  • The approach still relies on extensive pretraining and predictor accuracy, which may limit performance in unseen or highly specialized hardware environments.
  • While training costs are reduced, initial training remains resource-intensive, requiring substantial GPU hours and data collection.
  • Some sub-networks may still underperform in extreme resource constraints, necessitating further refinement of the architecture space and search strategies.

Future Work

Future directions include enhancing predictor accuracy, exploring more efficient search algorithms, and extending the framework to multi-task and multi-modal applications. Additionally, integrating automated meta-learning and reinforcement learning techniques could further reduce training costs and improve adaptation speed. Expanding the architecture space to include novel operations and hardware-specific optimizations will also be pursued to broaden applicability.

AI Executive Summary

Deep neural networks have revolutionized many AI applications, yet their increasing size and computational demands pose significant deployment challenges, especially on resource-limited devices. Traditional approaches—manual design, neural architecture search (NAS), and model compression—are either labor-intensive or computationally prohibitive, hindering scalability across diverse hardware platforms. To address this, the authors propose a novel framework called Once-for-All (OFA), which enables training a single, versatile network capable of supporting a vast architecture space with minimal additional cost.

The core innovation lies in a multi-dimensional progressive shrinking algorithm that supports elastic depth, width, kernel size, and input resolution. Starting from a large, full-capacity network, the method gradually fine-tunes smaller sub-networks by sharing weights and employing knowledge distillation, ensuring high accuracy across configurations. This process results in an extremely large architecture space (>10^19 configurations), all trained jointly, significantly reducing the need for repeated training for each deployment scenario.

Once trained, the OFA network can be rapidly specialized for specific hardware constraints using predictor-guided search, which estimates accuracy and latency without costly measurements. Extensive experiments on ImageNet demonstrate that OFA achieves 80.0% top-1 accuracy at less than 600M MACs, surpassing MobileNetV3 and EfficientNet in both accuracy and inference speed. The approach supports diverse hardware platforms—from GPUs to FPGAs and mobile phones—delivering speedups of 1.5 to 2.6 times over traditional models.

This work profoundly impacts scalable AI deployment, enabling a single training effort to serve many hardware and efficiency scenarios. It reduces training time, energy consumption, and carbon footprint, making high-performance models more accessible. Future work will focus on expanding the architecture space, improving predictor accuracy, and extending to multi-task settings, further broadening the framework’s applicability and efficiency.

Deep Analysis

Background

Deep learning模型在图像识别、自然语言处理等领域取得巨大成功,但模型规模不断扩大带来计算和能耗瓶颈。传统模型压缩技术如剪枝、量化和蒸馏已部分缓解问题,但难以满足多样化硬件平台的需求。神经架构搜索(NAS)通过自动设计网络结构提升效率,但训练成本高昂,难以在大规模场景中推广。近年来,面向硬件感知的NAS结合硬件反馈优化架构,但仍需为每个硬件平台重新搜索和训练,成本巨大。动态网络和多宽度支持模型如Slimmable Nets虽能在运行时调整结构,但受限于预定义架构,缺乏多维弹性。为解决这些问题,本文提出支持多维架构变化的统一训练框架,旨在实现模型的多场景适应性和高效部署。

Core Problem

现有方法在模型多样性和训练效率上存在明显瓶颈。手工设计和传统NAS需针对每个硬件平台重复训练,成本高昂且难以扩展。模型压缩虽能减小模型体积,但在多场景适配性方面不足。多宽度模型虽具弹性,但受限于预定义结构,难以满足复杂硬件需求。如何在保证模型性能的同时,实现一次训练支持多架构、多硬件环境的目标,成为关键难题。特别是在边缘设备、移动端等资源受限场景,需求更为迫切。解决方案需突破单一维度限制,支持深度、宽度、卷积核和输入分辨率的多维弹性,同时降低训练成本,提升部署效率。

Innovation

本研究的核心创新在于提出多维逐步收缩算法,支持深度、宽度、卷积核大小和输入分辨率的弹性调节,形成超百万级子网络空间。通过引入逐步收缩策略,避免不同子网络间的干扰,确保性能一致。采用预测器辅助架构搜索,实现快速匹配硬件平台和延迟约束,显著降低搜索成本。训练阶段采用共享权重和知识蒸馏技术,保证子网络性能。该方法实现了模型训练与搜索的解耦,极大降低了成本,为多场景部署提供了强大支撑。与传统单一架构或有限宽度模型相比,创新性地扩展了架构空间,支持多任务、多硬件环境的高效适配。

Methodology

  • �� 构建支持多维弹性架构的全局模型,定义深度、宽度、卷积核大小和输入分辨率的弹性空间。
  • �� 采用逐步收缩策略:先训练最大规模模型,然后逐步支持更小子网络,避免干扰。
  • �� 在训练中引入知识蒸馏,确保子网络性能。
  • �� 使用预测器指导架构搜索,根据硬件延迟和准确率快速筛选子网络。
  • �� 支持多平台多硬件环境的快速定制,解耦训练与搜索流程。
  • �� 通过多次实验验证不同配置的性能,优化模型参数和搜索策略。

Experiments

在ImageNet上,使用MobileNetV3架构空间,训练180 epochs,采用SGD优化。训练过程中应用逐步收缩算法,支持多维弹性架构。多平台(GPU、FPGA、手机)上测试,测量延迟和准确率。比较不同子网络性能,验证方法有效性。采用预测器和搜索算法,快速匹配硬件约束。实验还包括不同模型复杂度和架构配置的消融分析,验证训练效率和性能稳定性。

Results

OFA在ImageNet移动场景下达80.0%的Top-1准确率,模型复杂度仅595M MACs,优于MobileNetV3 4.0%的性能提升。训练成本显著降低,GPU小时数减少数十倍,碳排放降低数量级。多平台测试显示,支持多架构配置,延迟优化效果明显,速度提升1.5-2.6倍,准确率保持或超越独立训练模型。逐步收缩策略确保子网络在不同硬件环境下的优异表现,验证其广泛适应性。

Applications

该方法适用于移动端、边缘设备、云端多场景部署。可快速定制满足不同延迟、能耗、存储需求的模型,极大降低开发成本。未来可推广到多任务、多模态场景,实现动态架构调整,推动智能应用的普及。

Limitations & Outlook

当前方法仍依赖大量预训练和预测器,训练成本较高。模型压缩与架构搜索存在性能折中,部分子网络在极端资源限制下表现不足。未来需优化算法,降低训练资源消耗,增强泛化能力。

Plain Language Accessible to non-experts

想象你在准备一份多功能的厨房工具箱。传统上,你需要为不同的菜肴准备不同的厨具,比如炒锅、汤锅,每次都要换工具,既麻烦又浪费时间。现在,有了这个新方法,就像设计了一个万能厨具箱,里面的工具可以根据需要变成炒锅、汤锅甚至切菜板。只需一次准备,就能应对各种菜肴,无需每次都换工具。这个“万能厨具箱”就是论文中的一次性训练网络(OFA),它支持多种“架构配置”,可以根据不同硬件和需求“变形”。这样,无论是在家用手机、边缘设备还是云端服务器,只要选择合适的“工具”,就能快速得到高效的模型,既省时又省力。整个过程就像用一把变形魔法刀,随时变出你需要的工具,极大提高了效率和灵活性。

ELI14 Explained like you're 14

想象你有一个超级神奇的拼装玩具箱,里面的零件可以变成很多不同的玩具,比如汽车、飞机或者机器人。每次玩不同的游戏,你都不用买新玩具,只要用这些零件拼装就行了。以前,要做不同的玩具要花很多时间拆拆装装,还要重新设计每个玩具,特别麻烦。现在,这个神奇的玩具箱可以一次性装好很多不同的玩具,只要你告诉它你想玩什么,它就能变出对应的玩具来。这个神奇的“玩具箱”就像论文里的OFA网络,它可以根据不同的硬件设备和需求,快速变成最合适的模型。这样,无论是在手机上玩游戏,还是在电脑上做工作,都能用到最棒的工具,不用每次都重新做一遍。是不是很酷?这就像有个万能的变形玩具,随时变出你想要的东西,既方便又聪明!

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.

cs.LG cs.CV stat.ML