基础稀疏
昇腾模型压缩工具目前主要支持基于重训练的通道稀疏模型压缩特性,稀疏示例请参见获取更多样例>resnet101,支持通道稀疏的层以及约束如下:
优化方式 |
层名 |
约束 |
---|---|---|
通道稀疏 |
torch.nn.Linear:全连接层 |
depthwise只能被动稀疏(groups=in_channels),不能主动稀疏。 复用层(共用weight和bias参数)不支持稀疏。 |
torch.nn.Conv2d:卷积层 |
复用层(共用weight和bias参数)不支持稀疏。 |
接口调用流程
通道稀疏功能接口调用流程如图1所示。
图1 通道稀疏接口调用流程

蓝色部分为用户实现,灰色部分为用户调用昇腾模型压缩工具提供的API实现:
- 用户首先构造PyTorch的原始模型,调用create_prune_retrain_model接口对原始模型进行修改,修改后的模型参数量被裁剪。
- 对修改后的模型进行训练,直至精度满足要求;如果训练过程中断,则可基于原始模型和记录稀疏信息的文件,重新调用restore_prune_retrain_model接口稀疏模型,继续进行量化感知的训练,直至精度满足要求。
调用示例

- 如下示例标有“由用户补充处理”的步骤,需要用户根据自己的模型和数据集进行补充处理,示例中仅为示例代码。
- 调用昇腾模型压缩工具的部分,函数入参可以根据实际情况进行调整。稀疏基于用户的训练过程,请确保已经有基于PyTorch环境进行训练的脚本,并且训练后的精度正常。
- 导入昇腾模型压缩工具包,并通过安装后处理中的环境变量设置日志级别。
import amct_pytorch as amct
- (可选,由用户补充处理)建议使用原始待量化的模型和测试集,在PyTorch环境下推理,验证环境、推理脚本是否正常。
推荐执行该步骤,请确保原始模型可以完成推理且精度正常;执行该步骤时,可以使用部分测试集,减少运行时间。
ori_model.load() # 测试模型 user_test_model(ori_model, test_data, test_iterations)
- 调用昇腾模型压缩工具,执行通道稀疏流程。
- 稀疏模型。
实现该步骤前,应先恢复训练好的参数,如2中的ori_model.load()。
simple_cfg = './retrain.cfg' record_file = './tmp/record.txt' prune_retrain_model = amct.create_prune_retrain_model(model=ori_model, input_data=ori_model_input_data, config_defination=simple_cfg, record_file=record_file)
- (由用户补充处理)使用修改后的图,创建反向梯度,在训练集上做训练。
- 稀疏模型。
如果训练过程中断,需要从ckpt中恢复数据,继续训练,则调用流程为:
- 导入昇腾模型压缩工具包,并通过安装后处理中的环境变量设置日志级别。
import amct_pytorch as amct
- 准备原始模型。
ori_model= user_create_model()
- 调用昇腾模型压缩工具,恢复量化训练流程。
- 修改模型的参数和属性,减小模型的参数量,保存为新的prune_model。
simple_cfg = './retrain.cfg' record_file = './tmp/record.txt' prune_retrain_model = amct.restore_prune_retrain_model(model=ori_model, input_data=ori_model_input_data, record_file=record_file)
- (由用户补充处理)使用修改后的模型,恢复断点,创建反向梯度,在训练集上做训练。
- 修改模型的参数和属性,减小模型的参数量,保存为新的prune_model。
后续处理
经过稀疏后输出的模型为pth格式,该模型无法直接用于推理,需要用户自行将pth模型转成ONNX网络模型,然后才能使用ATC工具进行模型转换。
父主题: 通道稀疏