CANN/amct剪枝再训练模型创建

create_prune_retrain_model

【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct

产品支持情况

注：上述4选2结构化稀疏特性，标记“x”的产品，调用接口不会报错，但是获取不到性能收益。

功能说明

通道稀疏或4选2结构化稀疏接口，两种稀疏特性每次只能使能一个：将输入的待稀疏的图结构按照给定的稀疏配置文件进行稀疏处理，在传入的图结构中插入或者替换相关的算子，生成记录稀疏信息的record_file，返回修改后可用于稀疏后训练的torch.nn.Module模型。

函数原型

prune_retrain_model = create_prune_retrain_model (model, input_data, config_defination, record_file)

参数说明

通道稀疏支持的层及约束

表 1通道稀疏支持的层以及约束

4选2结构化稀疏支持的层及约束

表 1支持的层以及约束

返回值说明

返回修改后可用于稀疏后训练的torch.nn.Module模型。

调用示例

import amct_pytorch as amct # 建立待进行稀疏的网络图结构 model = build_model() model.load_state_dict(torch.load(state_dict_path)) input_data = tuple([torch.randn(input_shape)]) # 调用稀疏模型API record_file = os.path.join(TMP, 'scale_offset_record.txt') cfg_file = './prune_config.cfg' prune_retrain_model = amct.create_prune_retrain_model( model, input_data, cfg_file, record_file)

【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct