用于模型压缩的自动机器学习(AMC)

本仓库包含论文《AMC:用于移动设备上模型压缩和加速的自动机器学习》的PyTorch实现。

参考文献

如果您觉得本仓库有用,请引用我们的论文:

@inproceedings{he2018amc,
  title={AMC: AutoML for Model Compression and Acceleration on Mobile Devices},
  author={He, Yihui and Lin, Ji and Liu, Zhijian and Wang, Hanrui and Li, Li-Jia and Han, Song},
  booktitle={European Conference on Computer Vision (ECCV)},
  year={2018}
}

其他与自动化模型设计相关的论文:

• HAQ:硬件感知自动量化与混合精度(CVPR 2019)

• ProxylessNAS:直接在目标任务和硬件上进行神经架构搜索(ICLR 2019)

训练AMC

当前代码库支持在ImageNet上对MobileNet进行自动化剪枝。MobileNet的剪枝包括3个步骤:1.策略搜索;2.导出剪枝后的权重;3.从剪枝后的权重进行微调。

要执行完整的剪枝流程,请按以下说明操作(由于随机种子不同,结果可能与论文略有不同):

1. 策略搜索

   要在MobileNet ImageNet模型上搜索策略,首先通过运行以下命令获取在ImageNet上预训练的MobileNet检查点:

   bash ./checkpoints/download.sh
   

   它还将下载我们压缩到50% FLOPs的模型。然后运行以下脚本,在50% FLOPs约束下进行搜索:

   bash ./scripts/search_mobilenet_0.5flops.sh

   由于随机种子不同,结果可能有所不同。我们在论文中找到并报告的策略是:

   [3, 24, 48, 96, 80, 192, 200, 328, 352, 368, 360, 328, 400, 736, 752]
   

2. 导出剪枝后的权重

   搜索完成后,我们需要通过运行以下命令导出剪枝后的权重:

   bash ./scripts/export_mobilenet_0.5flops.sh
   

   我们还需要修改MobileNet文件以支持新的剪枝模型(这里已经在models/mobilenet.py中完成)

3. 从剪枝后的权重进行微调

   导出后,我们需要从剪枝后的权重进行微调。例如,我们可以通过运行以下命令使用余弦学习率进行150个epoch的微调:

   bash ./scripts/finetune_mobilenet_0.5flops.sh
   

AMC压缩模型

我们还提供了由AMC方法压缩的模型和权重。我们在这里提供了PyTorch和TensorFlow格式的压缩后的MobileNet-V1和MobileNet-V2。

详细统计数据如下:

依赖项

当前代码库在以下环境下进行了测试:

1. Python 3.7.3
2. PyTorch 1.1.0
3. torchvision 0.2.1
4. NumPy 1.14.3
5. SciPy 1.1.0
6. scikit-learn 0.19.1
7. tensorboardX
8. ImageNet数据集

联系方式

联系作者:

Ji Lin, jilin@mit.edu

Song Han, songhan@mit.edu