中文

仅训练一次:一种一次性神经网络训练与剪枝框架

机器学习 2021-11-15 v2

摘要

结构化剪枝是在资源受限设备上部署深度神经网络(DNNs)的常用技术。然而,现有剪枝方法通常是启发式的、任务特定的,并且需要额外的微调过程。为克服这些局限,我们提出了一种框架,通过“仅训练一次”(OTO)将 DNNs 压缩为具有竞争力性能且大幅降低 FLOPs 的更瘦架构。OTO 包含两个关键:(i)我们将 DNNs 的参数划分为零不变组,使我们能够剪枝零组而不影响输出;(ii)为促进零组,我们随后公式化一个结构化稀疏优化问题,并提出一种新颖优化算法——半空间随机投影梯度(HSPG)来求解,其在组稀疏探索上优于标准近端方法并保持可比收敛性。为展示 OTO 的有效性,我们从头同时训练并压缩完整模型而无需微调以加速推理并减少参数,在 CIFAR10 的 VGG16、CIFAR10 的 ResNet50 和 SQuAD 的 Bert 上取得最优结果,并在 ImageNet 的 ResNet50 上取得有竞争力结果。源代码可在 https://github.com/tianyic/only_train_once 获取。

关键词

引用

@article{arxiv.2107.07467,
  title  = {Only Train Once: A One-Shot Neural Network Training And Pruning Framework},
  author = {Tianyi Chen and Bo Ji and Tianyu Ding and Biyi Fang and Guanyi Wang and Zhihui Zhu and Luming Liang and Yixin Shi and Sheng Yi and Xiao Tu},
  journal= {arXiv preprint arXiv:2107.07467},
  year   = {2021}
}

备注

Accepted by NeurIPS 2021