[论文解读] AtomNAS: Fine-Grained End-to-End Neural Architecture Search

Jieru Mei, Yingwei Li+5 more|arXiv (Cornell University)|Dec 20, 2019

Advanced Neural Network Applications参考文献 33被引用 42

一句话总结

AtomNAS 引入了细粒度的原子块搜索空间和端到端可微分的 NAS，具备动态剪枝和 FLOPs 感知正则化，在降低搜索成本的同时实现了移动端 ImageNet 的最新结果。

ABSTRACT

Search space design is very critical to neural architecture search (NAS) algorithms. We propose a fine-grained search space comprised of atomic blocks, a minimal search unit that is much smaller than the ones used in recent NAS algorithms. This search space allows a mix of operations by composing different types of atomic blocks, while the search space in previous methods only allows homogeneous operations. Based on this search space, we propose a resource-aware architecture search framework which automatically assigns the computational resources (e.g., output channel numbers) for each operation by jointly considering the performance and the computational cost. In addition, to accelerate the search process, we propose a dynamic network shrinkage technique which prunes the atomic blocks with negligible influence on outputs on the fly. Instead of a search-and-retrain two-stage paradigm, our method simultaneously searches and trains the target architecture. Our method achieves state-of-the-art performance under several FLOPs configurations on ImageNet with a small searching cost. We open our entire codebase at: https://github.com/meijieru/AtomNAS.

研究动机与目标

• 需要比传统块级选择更细粒度的 NAS 搜索空间来提升需求的合理性。
• 提出原子块作为最小搜索单元，以实现混合操作和灵活的通道分配。
• 开发一个端到端的 NAS 框架，在训练网络权重的同时学习原子块的重要性，并执行 FLOPs 感知正则化。
• 引入动态网络收缩，通过在训练中剪除无用的原子块来加速搜索。
• 在移动 FLOP 约束下展示 ImageNet 的最先进性能，并证明对检测/分割的迁移性。

提出的方法

• 定义一个细粒度的搜索空间，即将一个构建块分解为由两个卷积通过通道维操作连接形成的原子块。
• 将整个网络表示为对原子块的求和，并引入可学习的重要性系数 alpha 实现端到端的块选择（式(4)）。
• 应用 FLOPs 感知正则化（对 alpha 乘以代价系数后的 L1 惩罚），以鼓励选择低成本的块（式(5)–(6)）。
• 使用训练中基于不断演化的 alpha 值的动态网络收缩来在训练期间移除死掉的原子块（算法1）。
• 剪枝后重新计算 BN 统计量，以确保推理性能的准确性。
• 使用端到端可微分的搜索与剪枝对超网进行端到端训练，最终得到的结构不需要额外的微调。

实验结果

研究问题

• RQ1 比传统的块级 NAS 在准确性与 FLOPs 的权衡上，细粒度的原子块搜索空间是否能获得更好结果？
• RQ2 在固定 FLOP 预算下，端到端可微分搜索结合 FLOPs 感知正则化，是否能够产生优于移动和紧凑设计基线的架构？
• RQ3 动态即时剪枝（收缩）是否能在不牺牲最终准确性的前提下降低搜索成本？
• RQ4 AtomNAS 生成的模型在目标检测和实例分割任务上的迁移效果如何？
• RQ5 剪枝后对 BN 重新标定对最终准确性的影响如何？

主要发现

• AtomNAS-C 在 ImageNet 上 360M FLOPs 时达到 75.9% 的 Top-1 准确率，超过若干最先进的移动模型。
• AtomNAS-C+ 在 363M FLOPs 时达到 77.6% 的 Top-1 准确率，超越 MixNet-M 并在 400M FLOPs 下确立移动场景的新 SOTA。
• AtomNAS-A 与 AtomNAS-B 在各自的 FLOPs 预算下也取得了出色结果（如 AtomNAS-A 在 258M FLOPs 下 74.6%，AtomNAS-B 在 326M FLOPs 下 75.5%）。
• 在消融实验中，与等惩罚基线相比，资源感知正则化改善了 FLOPs 效率与准确性之间的权衡。
• 剪枝后对 BN 的重新标定在 Top-1 准确率上带来显著提升（AtomNAS-A +1.4、AtomNAS-B +1.7、AtomNAS-C +1.2）。
• 动态收缩通过提前剪除无用块来降低搜索成本，AtomNAS-C 的总搜索+训练时间约为 25.5 小时（从头开始的最终训练约 22 小时）。
• AtomNAS 模型对 COCO 检测/分割的迁移性良好，在检测和分割指标上超越基线。