Once-for-All: Train One Network and Specialize it for Efficient Deployment
Han Cai1, Chuang Gan2, Tianzhe Wang1, Zhekai Zhang1, Song Han1 1Massachusetts Institute of Technology, 2MIT-IBM Watson AI Lab {hancai, chuangg,
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 Strubell et al. (2019)) 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
中文速览
面向海量IoT设备的深度学习部署面临一个核心难题：每换一块硬件或改变一次延迟约束，就得重新设计并从头训练一个专用网络，随着设备数量爆炸式增长，这种逐一训练的方式在算力和碳排放上代价极高。论文提出"一次训练，随处部署"的Once-for-All（OFA）网络：只训练一个超网络，推理时直接从中抽取符合目标硬件约束的子网络，无需任何额外训练，从而将专用模型设计的总成本从O(N)降至O(1)。为了让超过10¹⁹个共享权重的子网络都保持高精度，研究者设计了"渐进式收缩（progressive shrinking）"训练算法，先训好最大网络，再逐步将更小的子网络（涵盖深度、宽度、卷积核大小、输入分辨率四个维度）纳入优化，并配合知识蒸馏减少相互干扰。实验结果表明，OFA在CPU、手机GPU、FPGA等多种边缘设备上均超越MobileNetV3、EfficientNet等最先进方法（同精度下最高快2.6倍，或同速度下精度提升4%），并首次在移动端600M MACs约束下达到80%的ImageNet Top-1精度，同时将GPU计算开销和碳排放降低了数个数量级，对绿色高效的边缘AI部署具有重要意义。
原文 arXiv:1908.09791;中英对照 + 大白话阅读 https://aha.fim.ai/paper/1908.09791v5