Pruning-as-Search: Efficient Neural Architecture Search via Channel Pruning and Structural Reparameterization
Yanyu Li1111Work is done at Kwai Inc. Pu Zhao1 Geng Yuan1 Xue Lin1 Yanzhi Wang1 Xin Chen 1Northeastern University 2Intel Corp. {li.yanyu, zhao.pu, yuan.geng, xue.lin,
Abstract
Neural architecture search (NAS) and network pruning are widely studied efficient AI techniques, but not yet perfect. NAS performs exhaustive candidate architecture search, incurring tremendous search cost. Though (structured) pruning can simply shrink model dimension, it remains unclear how to decide the per-layer sparsity automatically and optimally. In this work, we revisit the problem of layer-width optimization and propose Pruning-as-Search (PaS), an end-to-end channel pruning method to search out desired sub-network automatically and efficiently. Specifically, we add a depth-wise binary convolution to learn pruning policies directly through gradient descent. By combining the structural reparameterization and PaS, we successfully searched out a new family of VGG-like and lightweight networks, which enable the flexibility of arbitrary width with respect to each layer instead of each stage. Experimental results show that our proposed architecture outperforms prior arts by around $1.0\%$ top-1 accuracy under similar inference speed on ImageNet-1000 classification task. Furthermore, we demonstrate the effectiveness of our width search on complex tasks including instance segmentati
中文速览
自动化神经网络宽度设计长期面临两大困境:NAS搜索代价极高,而传统剪枝又难以自动、合理地决定每层的稀疏度。本文提出"剪枝即搜索"(Pruning-as-Search, PaS)方法,在每个卷积层后插入一个深度可分离二值卷积层作为剪枝指示器,借助直通估计器(Straight Through Estimator, STE)实现端到端梯度优化,从而让网络在一次微调过程中自动学习每层的剪枝策略,彻底摆脱了依赖权重幅度判断重要性所带来的"幅度陷阱"。此外,论文将PaS与结构重参数化(structural reparameterization)结合,训练时保留残差连接以稳定优化,推理时将其融合为普通卷积,从而突破了"同一阶段各层宽度必须相同"的刚性约束,实现了逐层任意宽度的设计自由度,并由此衍生出一族全新的VGG式及轻量级网络。在ImageNet-1000分类任务上,所搜索的架构在相近推理速度下比现有方法高出约1%的Top-1准确率,同时在实例分割和图像翻译等复杂任务上也验证了方法的有效性。
原文 arXiv:2206.01198;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2206.01198v1