[论文解读] Universally Slimmable Networks and Improved Training Techniques

Jiahui Yu, Thomas S. Huang|arXiv (Cornell University)|Mar 12, 2019

Advanced Neural Network Applications参考文献 23被引用 16

一句话总结

本文提出通用剪枝网络（US-Nets），一种训练单个神经网络的方法，使其能够在连续宽度范围内高效推理，而不仅限于预定义的离散宽度。通过引入夹心规则（sandwich rule）和就地蒸馏（inplace distillation），US-Nets 在所有宽度下的准确率均优于单独训练的模型或离散剪枝网络，且推理和训练过程的计算开销极低。

ABSTRACT

Slimmable networks are a family of neural networks that can instantly adjust the runtime width. The width can be chosen from a predefined widths set to adaptively optimize accuracy-efficiency trade-offs at runtime. In this work, we propose a systematic approach to train universally slimmable networks (US-Nets), extending slimmable networks to execute at arbitrary width, and generalizing to networks both with and without batch normalization layers. We further propose two improved training techniques for US-Nets, named the sandwich rule and inplace distillation, to enhance training process and boost testing accuracy. We show improved performance of universally slimmable MobileNet v1 and MobileNet v2 on ImageNet classification task, compared with individually trained ones and 4-switch slimmable network baselines. We also evaluate the proposed US-Nets and improved training techniques on tasks of image super-resolution and deep reinforcement learning. Extensive ablation experiments on these representative tasks demonstrate the effectiveness of our proposed methods. Our discovery opens up the possibility to directly evaluate FLOPs-Accuracy spectrum of network architectures. Code and models are available at: https://github.com/JiahuiYu/slimmable_networks

研究动机与目标

• 通过将剪枝网络扩展至连续宽度谱，实现神经网络在任意宽度下运行，而不仅限于离散宽度。
• 解决在批量归一化（batch normalization）下训练通用剪枝网络的挑战，传统方法因训练-推理统计不一致而失效。
• 通过新颖的训练技术，提升通用剪枝网络的训练效率和测试准确率。
• 实现对单个模型的 FLOPs-准确率权衡谱的直接评估，无需训练多个独立模型。

提出的方法

• 提出一种训练后批量归一化统计计算方法：在模型权重固定后，对一个小样本集（如 1,024 张图像）的批量统计量计算移动平均或精确平均，实现快速且准确的推理。
• 引入夹心规则：在训练过程中，每轮迭代采样多个宽度（如 4 个宽度），对每个宽度计算损失，并通过加权平均组合，以稳定训练并提升泛化能力。
• 开发就地蒸馏技术：在训练过程中，将更宽子网络的知识蒸馏到更窄子网络中，从而在不增加推理成本的前提下，提升所有宽度下的准确率。
• 通过避免在训练期间累积各宽度的统计量，对批量归一化进行改进；相反，使用一个小而具有代表性的数据子集在训练后统一计算所有 BN 统计量。
• 采用有界残差误差公式，从理论上证明更宽的网络不应表现差于更窄的网络，从而支持通用剪枝的可行性。
• 使用宽度除数（d=8）将通道数对齐至硬件对齐要求（如 GPU warp size），提升推理速度，并确保与 MobileNet 的 FLOPs 比较公平。

实验结果

研究问题

• RQ1 能否训练单个神经网络，使其在连续宽度范围内均表现准确，而不仅限于离散宽度？
• RQ2 如何在通用剪枝网络中有效处理批量归一化，以确保所有宽度下的准确率？
• RQ3 哪些训练策略可使通用剪枝网络的性能优于单独训练的模型或离散剪枝基线？
• RQ4 能否使用单个模型直接评估网络的 FLOPs-准确率权衡谱？这对模型部署有何影响？
• RQ5 超参数如宽度下界、宽度除数和每轮迭代采样宽度数如何影响 US-Nets 的性能？

主要发现

• 采用夹心规则和就地蒸馏训练的 US-Nets 在 ImageNet 上实现 28.2% 的 Top-1 错误率，使用单个模型，优于所有宽度下的单独训练的 MobileNet v1 模型和 4-切换剪枝基线。
• 仅使用 1,024 张图像计算的训练后批量归一化统计量，性能几乎与全批量统计量相同，实现快速且准确的推理。
• US-Net 的性能从根本上受限于其最小宽度（k₀），从 0.25× 到 1.0× 训练的模型，其准确率优于最小宽度更高的模型（如 0.35× 或 0.05×）。
• 使用宽度除数 d=8 可略微提升性能，并确保硬件友好的通道数，与 d=1 相比准确率无显著下降。
• 每轮迭代采样 4 个宽度的训练策略在准确率与效率之间取得最佳平衡，采样 5 个宽度仅带来边际收益，而采样 3 个宽度则表现较差。
• 就地蒸馏显著提升了所有宽度下的测试准确率，且无额外推理成本，表明其是训练鲁棒剪枝网络的关键组件。