用于模型压缩的自动机器学习(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}
}
其他与自动化模型设计相关的论文:
训练AMC
当前代码库支持在ImageNet上对MobileNet进行自动化剪枝。MobileNet的剪枝包括3个步骤:1.策略搜索;2.导出剪枝后的权重;3.从剪枝后的权重进行微调。
要执行完整的剪枝流程,请按以下说明操作(由于随机种子不同,结果可能与论文略有不同):
-
策略搜索
要在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]
-
导出剪枝后的权重
搜索完成后,我们需要通过运行以下命令导出剪枝后的权重:
bash ./scripts/export_mobilenet_0.5flops.sh
我们还需要修改MobileNet文件以支持新的剪枝模型(这里已经在
models/mobilenet.py
中完成) -
从剪枝后的权重进行微调
导出后,我们需要从剪枝后的权重进行微调。例如,我们可以通过运行以下命令使用余弦学习率进行150个epoch的微调:
bash ./scripts/finetune_mobilenet_0.5flops.sh
AMC压缩模型
我们还提供了由AMC方法压缩的模型和权重。我们在这里提供了PyTorch和TensorFlow格式的压缩后的MobileNet-V1和MobileNet-V2。
详细统计数据如下:
模型 | Top1 准确率 (%) | Top5 准确率 (%) |
---|---|---|
MobileNetV1-width*0.75 | 68.4 | 88.2 |
MobileNetV1-50%FLOPs | 70.494 | 89.306 |
MobileNetV1-50%Time | 70.200 | 89.430 |
MobileNetV2-width*0.75 | 69.8 | 89.6 |
MobileNetV2-70%FLOPs | 70.854 | 89.914 |
依赖项
当前代码库在以下环境下进行了测试:
- Python 3.7.3
- PyTorch 1.1.0
- torchvision 0.2.1
- NumPy 1.14.3
- SciPy 1.1.0
- scikit-learn 0.19.1
- tensorboardX
- ImageNet数据集
联系方式
联系作者:
Ji Lin, jilin@mit.edu
Song Han, songhan@mit.edu