CANN AMCT create_compressed_retrain_model 接口使用指南稀疏与量化的静态组合压缩【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct导读create_compressed_retrain_model是 CANN AMCTAscend Model Compression Toolkit昇腾 AI 处理器亲和模型压缩工具提供的静态组合压缩接口用于将稀疏通道稀疏或 4 选 2 结构化稀疏与量化感知训练QAT两个压缩手段在一个模型上串联执行先稀疏、后量化最终返回可直接进行重训练的torch.nn.Module模型。读完本文你将掌握该接口的函数原型、四个入参的语义与类型、组合压缩配置文件的完整编写方法含可直接复用的配置样例、调用示例与落盘文件行为并能结合仓库源码理解接口内部解析图 → 稀疏改图 → 量化插层的真实执行链路。产品支持情况该接口在不同昇腾产品上的特性支持情况如下表所示产品量化感知训练通道稀疏4选2结构化稀疏Ascend 950PR / Ascend 950DTINT8 量化√INT4 量化x√xAtlas A3 训练系列产品 / Atlas A3 推理系列产品INT8 量化√INT4 量化x√√Atlas A2 训练系列产品 / Atlas A2 推理系列产品INT8 量化√INT4 量化x√√注特性中标记x的产品调用接口不会报错但是获取不到性能收益。例如在 Ascend 950PR/Ascend 950DT 上配置 4 选 2 结构化稀疏可以正常执行压缩流程但无法获得该特性带来的硬件加速收益。功能说明先稀疏后量化的静态组合压缩静态组合压缩接口的职责是将输入的待静态组合压缩的模型按照给定的组合压缩配置文件进行压缩处理具体处理流程为稀疏对模型执行通道稀疏或 4 选 2 结构化稀疏两种稀疏特性二选一量化在稀疏后的模型上插入量化相关的算子包括数据和权重的量化感知训练层以及 searchN 的层记录生成稀疏和量化因子记录文件record_file如果配置存在返回返回修改后的torch.nn.Module模型供后续量化感知训练使用。从仓库源码 prune_interface.py 中的_modify_original_to_compressed_model可以看到接口内部的真实执行步骤这与文档描述完全对应step 1 解析通过Parser.export_onnx(model, input_data, model_onnx)将 PyTorch 模型导出为 ONNX 图再调用Parser.parse_net_to_graph解析为内部图结构graph.add_model(model)绑定原始模型随后调用RetrainConfig.init(graph, config_defination, enable_retrainTrue, enable_pruneTrue)解析组合压缩配置同时使能量化重训练与稀疏两个特性。step 2 稀疏当RetrainConfig.enable_prune为真时构建PruneHelper(graph, input_data, record_file)并调用create_prune_model()执行通道稀疏filter prune随后调用create_selective_prune_record(graph)生成 4 选 2 结构化稀疏记录。step 3 稀疏改图重新导出并解析图后调用_modify_original_model_to_prune(model, graph)将稀疏结果真正改写回原始模型中。step 4 量化插层当RetrainConfig.enable_retrain为真时调用_modify_original_model_to_quant(model, record_file, graph)插入数据量化ULQ、权重量化ARQ等量化感知训练层与 searchN 层并在前向推理过程中将量化因子写入record_file。入口函数本身prune_interface.py还会执行参数检查check_params校验model必须是torch.nn.Module、config_defination与record_file必须是str、ModuleHelper.check_amct_op()检查、ModuleHelper.deep_copy(model)深拷贝并通过files_util.create_empty_file(record_file, check_existTrue)创建/清空记录文件通过SingletonScaleOffsetRecord().reset_singleton(record_file)重置量化因子单例记录器。这意味着接口不会原地修改你传入的原始模型返回的是深拷贝后改造的模型便于对比与回滚。函数原型compressed_retrain_model create_compressed_retrain_model(model, input_data, config_defination, record_file)参数说明参数名输入/输出说明model输入含义PyTorch 的 model。数据类型torch.nn.Moduleinput_data输入含义模型的输入数据用于模型导出与图解析可为随机数据。数据类型tupleconfig_defination输入含义静态组合压缩简易配置文件。说明基于retrain_config_pytorch.proto文件生成的简易配置文件compressed.cfg*.proto文件所在路径为AMCT安装目录/amct_pytorch/proto/。*.proto文件参数解释以及生成的compressed.cfg简易配置文件样例请参见 量化感知训练简易配置文件。数据类型stringrecord_file输入含义待记录稀疏和量化因子文件路径及名称。数据类型string说明config_defination与record_file均为文件路径字符串。参考 接口列表与整体约束AMCT 不会对用户输入的路径做安全校验请确保路径正确若重复执行量化同一路径下的文件内容会被覆盖打屏日志中会有相关文件被覆盖的 warning 风险提示。返回值说明返回根据配置文件完成处理后的模型已按配置进行稀疏如果配置了稀疏且插入量化相关层如果配置了量化的torch.nn.Module静态组合压缩模型可直接用于量化感知训练。约束说明组合压缩配置文件至少存在一个配置稀疏配置或者量化配置。即prune_config与量化配置retrain_data_quant_config/retrain_weight_quant_config二者必须至少出现其一否则接口无实际压缩动作。通道稀疏与 4 选 2 结构化稀疏两种稀疏特性每次只能使能一个。组合压缩简易配置文件的编写config_defination指向的组合压缩配置文件由AMCT安装目录/amct_pytorch/proto/retrain_config_pytorch.proto仓库中对应源文件为 retrain_config_pytorch.proto定义。该 proto 可同时生成量化感知训练简易配置文件、稀疏简易配置文件、组合压缩简易配置文件三种类型组合压缩场景只需在同一个.cfg中同时写入稀疏配置与量化配置。核心配置项如下prune_config稀疏配置全局内部可选filter_pruner.balanced_l2_norm_filter_prune通道稀疏BCP 算法或n_out_of_m_pruner.l1_selective_prune4 选 2 结构化稀疏。balanced_l2_norm_filter_prune.prune_ratio必填float稀疏率被稀疏的 filter 数量与 filter 总数量的比值推荐配置为 0.2即裁剪 20% 的输出通道balanced_l2_norm_filter_prune.ascend_optimized可选bool是否做昇腾亲和优化若稀疏后的模型要部署在 AI 处理器上建议配置为 truel1_selective_prune.n_out_of_m_type可选目前仅支持M4N2即每 4 个连续权重中保留 2 个权重l1_selective_prune.update_freq可选uint32更新 4 选 2 稀疏选择的间隔update_freq0时仅在第一个 batch 更新update_freq2时每 2 个 batch 更新一次默认为 0。retrain_data_quant_config数据量化全局配置ulq_quantizeULQ 数据量化算法内可配置dst_type数据量化位宽仅支持 INT8、clip_max_min.clip_max/clip_min初始上下限不配置则默认用 IFMR 初始化、fixed_min是否下限固定为 0。retrain_weight_quant_config权重量化全局配置arq_retrainARQ 权重量化算法内可配置dst_typeINT8默认与channel_wise是否做 channel wise 的 ARQ。skip_layers/skip_layer_types按层名 / 按层类型跳过哪些层不做压缩全局参数稀疏场景专用跳过参数为regular_prune_skip_layers/regular_prune_skip_types量化场景专用跳过参数为quant_skip_layers/quant_skip_types。若同时配置全局与场景专用参数则取两者并集。override_layer_configs/override_layer_types按层名 / 按层类型对指定层做差异化压缩如差异化稀疏率或差异化量化位宽。参数优先级为override_layer_configs override_layer_types prune_config / retrain_data_quant_config / retrain_weight_quant_config。组合压缩通道稀疏 INT8 量化配置样例以下_compressed1_.cfg样例展示了通道稀疏与 INT8 量化的组合配置可直接作为模板改写详见 量化感知训练简易配置文件prune_config : { filter_pruner : { balanced_l2_norm_filter_prune : { prune_ratio : 0.3 ascend_optimized: True } } } # skip_layers: skip_layers_name_0 skip_layer_types: Optype quant_skip_layers: Opname quant_skip_types: Optype retrain_weight_quant_config: { arq_retrain: { channel_wise: true dst_type: INT8 } } override_layer_types : { layer_type: Optype retrain_weight_quant_config: { arq_retrain: { channel_wise: false dst_type: INT8 } } retrain_data_quant_config : { ulq_quantize : { clip_max_min : { clip_max : 6.0 clip_min : -6.0 } } } prune_config : { filter_pruner : { balanced_l2_norm_filter_prune : { prune_ratio : 0.5 ascend_optimized: True } } } }组合压缩4选2结构化稀疏 INT8 量化配置样例以下_compressed2_.cfg样例展示了 4 选 2 结构化稀疏与 INT8 量化的组合配置prune_config{ n_out_of_m_pruner { l1_selective_prune { n_out_of_m_type: M4N2 update_freq: 0 } } } # skip_layers: skip_layers_name_0 skip_layer_types: Optype quant_skip_layers: quant_skip_layers_name_0 quant_skip_types: Optype retrain_weight_quant_config: { arq_retrain: { channel_wise: true dst_type: INT8 } } override_layer_types : { layer_type: Optype retrain_weight_quant_config: { arq_retrain: { channel_wise: false dst_type: INT8 } } retrain_data_quant_config : { ulq_quantize : { clip_max_min : { clip_max : 6.0 clip_min : -6.0 } } } prune_config{ n_out_of_m_pruner { l1_selective_prune { n_out_of_m_type: M4N2 update_freq: 1 } } } }需要特别注意的是由于硬件约束Ascend 950PR / Ascend 950DT 不支持 4 选 2 结构化稀疏特性在这些产品上请使用通道稀疏方案。调用示例以下代码展示了create_compressed_retrain_model的完整调用流程build_model为示例中的待压缩网络构建函数input_shape为模型输入张量形状import os import torch import amct_pytorch as amct # 建立待进行静态组合压缩的网络 model build_model() input_data tuple([torch.randn(input_shape)]) # 调用静态组合压缩API TMP ./tmp record_file os.path.join(TMP, compressed_record.txt) config_defination ./compressed_cfg.cfg compressed_retrain_model amct.create_compressed_retrain_model( model, input_data, config_defination, record_file)其中input_data传入tuple([torch.randn(input_shape)])即可——源码实现仅用它完成 ONNX 导出与图解析Parser.export_onnx注释中也明确说明 input_data: used to compile model, can be random data因此可以使用随机数据。record_file会在接口内部被创建/清空为空的记录文件files_util.create_empty_file(record_file, check_existTrue)无需预先创建。compressed_retrain_model即为改图后的静态组合压缩训练模型可直接进入量化感知训练流程。落盘文件说明接口保存的静态组合压缩记录文件record_file如果简易配置文件中含有稀疏配置则在该函数完成后record_file中含有稀疏记录信息包括 filter 稀疏的选择结果以及 4 选 2 结构化稀疏的记录。量化因子则在前向传播训练/推理过程中由插入的量化层计算并写入同一记录文件。该文件后续会作为 restore_compressed_retrain_model 的输入用于恢复压缩训练模型并加载 checkpoint 权重最终通过 save_compressed_retrain_model 生成部署模型与精度仿真模型。完整工作流与测试验证create_compressed_retrain_model属于 接口列表 中组合压缩接口的入口之一与restore_compressed_retrain_model、save_compressed_retrain_model共同构成压缩训练 → 恢复训练 → 导出部署的完整链路create_compressed_retrain_model稀疏 量化插层产出可训练模型与record_file用户对返回模型执行量化感知训练保存 checkpointrestore_compressed_retrain_model按record_file稀疏记录重建模型、加载训练权重save_compressed_retrain_model导出 ONNX 仿真模型与部署模型。从仓库测试用例 test_prune_interface.py 可以看到该系列接口的标准用法以随机输入构建模型 → 调用create_prune_retrain_model/create_compressed_retrain_model得到新模型 → 校验前后向输出 shape 一致 →torch.save({state_dict: new_model.state_dict()}, pth_file)保存权重 → 用restore_*系列接口恢复并校验输出完全一致(new_output new_output2).all()为 True。这套模式同样适用于组合压缩接口的验证稀疏与量化插层不应改变模型的输入输出 shape且恢复接口必须能精确还原训练中的模型状态。总结create_compressed_retrain_model以一份.cfg配置同时驱动稀疏 量化两种压缩特性在模型内部形成先稀疏、后量化的静态组合压缩图。使用时的关键要点是配置文件必须至少包含稀疏或量化配置之一稀疏的两种方式通道稀疏 / 4 选 2 结构化稀疏只能二选一record_file是串联后续恢复与导出流程的纽带。结合仓库源码prune_interface.py、retrain_config_pytorch.proto与测试用例你可以完整把握从配置编写、接口调用到训练恢复、模型导出的全流程将稀疏与量化感知训练组合应用于昇腾 AI 处理器上的模型压缩部署。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考