AMCT create_distill_config 蒸馏量化配置生成接口详解原理、参数与配置格式【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct导读本文围绕 CANN AMCT 开源仓amct_pytorch提供的create_distill_config接口展开说明如何基于已加载权重的浮点模型与输入数据自动扫描图结构中的可蒸馏量化层与蒸馏结构并落盘生成 JSON 格式的蒸馏量化配置文件。读完本文你将掌握该接口的函数原型、全部参数约束、内部实现链路ONNX 导出 → 图解析 → 配置生成以及配套蒸馏简易配置文件distill.cfg/distill_config_pytorch.proto的完整字段语义可直接用于蒸馏量化训练的前置配置阶段。产品支持情况该接口在以下昇腾产品上获得支持依据 create_distill_config.md产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√功能说明create_distill_config是 AMCT 的蒸馏配置生成接口它根据输入模型的图结构自动找出所有可蒸馏量化的层与可蒸馏量化的结构自动生成蒸馏量化配置文件并将可蒸馏量化层的量化配置和蒸馏结构一并写入该配置文件。换句话说它把人工梳理网络、手写每层量化配置与蒸馏分组的工作自动化你只需提供模型与一份输入样例接口即可产出后续create_distill_model、distill蒸馏流程所依赖的config.json。函数原型create_distill_config(config_file, model, input_data, config_definationNone)参数说明参数名输入/输出说明config_file输入含义待生成的蒸馏量化配置文件存放路径及名称。如果存放路径下已经存在该文件则调用该接口时会覆盖已有文件。数据类型stringmodel输入含义待进行蒸馏量化的原始浮点模型已加载权重。数据类型torch.nn.Moduleinput_data输入含义模型的输入数据。一个 torch.tensor 会被等价为 tuple(torch.tensor)。数据类型tupleconfig_defination输入含义简易配置文件基于distill_config_pytorch.proto文件生成的简易配置文件 distill.cfg*.proto文件所在路径为 AMCT 安装目录/amct_pytorch/proto/。默认值None。数据类型string关于config_defination参数有两个关键行为可在 distillation_interface.py 的源码中得到印证当config_defination为None时接口调用create_default_distill_config(config_file, graph)即完全依据图结构生成默认蒸馏配置当传入简易配置文件路径时接口先校验文件存在性不存在则抛出FileNotFoundError再调用create_distill_config_from_proto(config_file, graph, config_proto_file)将 proto 格式的简易配置展开为完整 JSON 配置。此外源码还体现了两个内部细节如果传入的模型是torch.nn.parallel.DistributedDataParallel包装对象接口会自动取出.module再进行解析模型先经Parser.export_onnx(model, input_data, tmp_onnx)导出为 ONNX再经Parser.parse_net_to_graph(tmp_onnx)解析为 IR 图最终graph.add_model(model)绑定原始模型后进入配置生成逻辑——这正是根据图的结构查找可蒸馏量化层的实现基础。返回值说明无。调用示例import amct_pytorch as amct # 建立待进行蒸馏量化的网络图结构 model build_model() model.load_state_dict(torch.load(state_dict_path)) input_data tuple([torch.randn(input_shape)]) # 生成蒸馏配置文件 amct.create_distill_config(config_file./configs/config.json, model, input_data, config_defination./configs/distill.cfg)落盘文件说明接口生成 JSON 格式的蒸馏量化配置文件。重新执行蒸馏时该接口输出的配置文件将会被覆盖。以下为 INT8 量化场景下的配置文件样例{ version:1, batch_num:1, group_size:1, data_dump:false, distill_group:[ [ conv1, bn, relu ], [ conv2, bn2, relu2 ] ], conv1:{ quant_enable:true, distill_data_config:{ algo:ulq_quantize, dst_type:INT8 }, distill_weight_config:{ algo:arq_distill, channel_wise:true, dst_type:INT8 } }, conv2:{ quant_enable:true, distill_data_config:{ algo:ulq_quantize, dst_type:INT8 }, distill_weight_config:{ algo:arq_distill, channel_wise:true, dst_type:INT8 } } }从仓库的测试用例可看到与实际运行完全一致的落盘结构tests/amct_pytorch/testcase_python/distill/cfgs/distill_cfg.json中即包含version、batch_num、group_size、data_dump、distill_group与逐层的distill_data_config、distill_weight_config字段并被test_distill_interface.pytests/amct_pytorch/testcase_python/distill/test_distill_interface.py中test_create_distill_config_no_cfg、test_create_distill_config_cfg_exist等用例用于验证配置文件的生成与覆盖行为。JSON 配置顶层字段说明结合 distill_config.py 与落盘样例生成的 JSON 文件包含以下核心要素version配置文件版本号当前为 1batch_num蒸馏 batch 数量用于 ifmr 积累数据计算量化因子group_size蒸馏 block 中最小蒸馏单元个数data_dumpteacher 网络 block 输入输出 dump 开关distill_group蒸馏结构分组每组为一个层名列表如[conv1,bn,relu]表示这些层作为一个蒸馏单元整体参与蒸馏逐层配置以层名为 key包含quant_enable该层是否量化、distill_data_config数据量化配置algo为ulq_quantize时配置dst_type、distill_weight_config权重量化配置algo为arq_distill或ulq_distill时配置channel_wise与dst_type。蒸馏简易配置文件distill.cfg与 proto 定义config_defination参数传入的是基于distill_config_pytorch.proto生成的简易配置文件。该 proto 文件在仓库中的实际位置为 amct_pytorch/classic/graph_based/amct_pytorch/proto/distill_config_pytorch.proto完整字段说明详见 蒸馏简易配置文件。简易配置文件的完整字段语义如下表所示消息是否必填类型字段说明AMCTDistillConfig---AMCT 蒸馏的简易配置optionaluint32batch_num蒸馏 batch 数量用于 ifmr 积累数据计算量化因子optionaluint32group_size蒸馏 block 中最小蒸馏单元个数optionalbooldata_dumpteacher 网络 block 输入输出 dump 开关repeatedDistillGroupdistill_group用户自定义蒸馏结构optionalDistillDataQuantConfigdistill_data_quant_config蒸馏数据量化配置参数optionalDistillWeightQuantConfigdistill_weight_quant_config蒸馏权重量化配置参数repeatedDistillOverrideLayerdistill_override_layers重写某一层的量化配置repeatedDistillOverrideLayerTypedistill_override_layer_types重写某一类型层的量化配置repeatedstringquant_skip_layers不需要量化的层仍然做蒸馏repeatedstringquant_skip_layer_types不需要量化的算子类型仍然需要做蒸馏DistillGroup---用户自定义蒸馏结构。蒸馏结构中仅支持 torch.nn.Module 类型的算子requiredstringstart_layer_name用户自定义蒸馏结构起始层requiredstringend_layer_name用户自定义蒸馏结构结束层DistillDataQuantConfig---蒸馏数据量化配置-ActULQquantizeulq_quantize数据量化的算法目前仅支持 ulqActULQquantize---ULQ 数据量化算法配置算法介绍见 ULQ 数据量化算法optionalClipMaxMinclip_max_min初始化的上下限值如果不配置默认用 ifmr 进行初始化optionalboolfixed_min是否下限不学习且固定为 0。默认 ReLU 后为 true其他为 falseoptionalDataTypedst_type用以选择 INT8 或 INT4 量化位宽默认为 INT8。当前版本仅支持 INT8 量化ClipMaxMin---初始上下限requiredfloatclip_max初始上限值requiredfloatclip_min初始下限值DistillWeightQuantConfig---蒸馏权重量化配置-ARQDistillarq_distillARQ 权重量化算法-WtsULQDistillulq_distillULQ 权重量化算法ARQDistill---ARQ 权重量化算法配置算法介绍见 ARQ 权重量化算法optionalDataTypedst_type用以选择 INT8 或 INT4 量化位宽默认为 INT8。当前版本仅支持 INT8 量化optionalboolchannel_wise是否做 channel wise 的 arqWtsULQDistill---ULQ 权重量化算法配置算法介绍见 ULQ 数据量化算法optionalDataTypedst_type用以选择 INT8 或 INT4 量化位宽默认为 INT8。当前版本仅支持 INT8 量化optionalboolchannel_wise是否做 channel wise 的 ulqDistillOverrideLayer---重写的层配置requiredstringlayer_name层名optionalDistillDataQuantConfigdistill_data_quant_config重写的数据层量化参数optionalDistillWeightQuantConfigdistill_weight_quant_config重写的权重层量化参数DistillOverrideLayerType---重写的层类型配置requiredstringlayer_type层类型optionalDistillDataQuantConfigdistill_data_quant_config重写的数据层量化参数optionalDistillWeightQuantConfigdistill_weight_quant_config重写的权重层量化参数简易配置文件完整样例基于上述 proto 定义构造的蒸馏简易配置文件quant.cfg样例如下可直接对照 蒸馏简易配置文件batch_num: 1 group_size: 1 data_dump: true distill_group: { start_layer_name: layer1 end_layer_name: layer2 } distill_data_quant_config: { ulq_quantize: { clip_max_min: { clip_max: 6.0 clip_min: -6.0 } fixed_min: true dst_type: INT8 } } distill_weight_quant_config: { arq_distill: { channel_wise: true dst_type: INT8 } } quant_skip_layers: layer3 quant_skip_layer_types: type1 distill_override_layers : { layer_name: layer4 distill_data_quant_config: { ulq_quantize: { clip_max_min: { clip_max: 3.0 clip_min: -3.0 } fixed_min: true dst_type: INT8 } } distill_weight_quant_config: { arq_distill: { channel_wise: false dst_type: INT8 } } } distill_override_layer_types : { layer_type: type2 distill_data_quant_config: { ulq_quantize: { clip_max_min: { clip_max: 3.0 clip_min: -3.0 } fixed_min: true dst_type: INT8 } } distill_weight_quant_config: { ulq_distill: { channel_wise: false dst_type: INT8 } } }简易配置文件的常见用法从仓库测试用例 cfgs/config.cfg 可见最简单的合法配置可以只包含一行data_dump: true该用例通过create_distill_config(config_file, model, input_data, config_defination./cfgs/config.cfg)验证了简易配置 图结构可以成功生成完整 JSON 配置test_create_distill_config_cfg_not_exist则验证了当配置文件不存在时会抛出FileNotFoundError。这说明简易配置文件的主要作用是对默认行为做定向覆盖通过data_dump开启/关闭 teacher 网络 block 输入输出 dump通过distill_data_quant_config/distill_weight_quant_config设置全局默认的数据/权重量化算法、clip_max_min初始化范围、fixed_min与channel_wise通过distill_override_layers/distill_override_layer_types按层名或层类型局部重写量化参数通过quant_skip_layers/quant_skip_layer_types指定不量化但参与蒸馏的层或算子类型。生成配置之后的蒸馏流程衔接create_distill_config生成的config.json是蒸馏量化的第一步其后续调用链在 distillation_interface.py 中均有实现create_distill_model(config_file, model, input_data)读取该 JSON 配置通过parse_distill_config解析并执行InsertQatPass将普通层替换为可训练的量化模块返回修改后的蒸馏模型学生模型distill(model, compress_model, config_file, train_loader, epochs1, lr1e-3, sample_instanceNone, lossNone, optimizerNone)以原始模型为 teacher、压缩模型为 student按配置中的distill_group分组逐一进行蒸馏训练返回蒸馏后的压缩模型save_distill_model(model, save_path, input_data, ...)蒸馏完成后导出fake_quant与deploy两个 ONNX 模型。其中distill接口要求传入的config.json必须与create_distill_config接口产出的config.json保持一致即先生成配置、再创建蒸馏模型、最后执行蒸馏是完整的标准流程。相关接口说明可参见 create_distill_model 与 distill。使用注意事项config_file已存在时会被覆盖请避免在蒸馏中途复用已被修改的配置文件简易配置文件config_defination为可选参数不传时接口完全依据图结构生成默认配置distill_group仅支持 torch.nn.Module 类型的算子参与蒸馏结构定义当前版本的数据量化ulq_quantize与权重量化arq_distill/ulq_distill仅支持 INT8 量化位宽dst_type默认 INT8INT4 枚举值在 proto 中已定义但当前版本未开放input_data支持单个torch.tensor或tuple单个 tensor 会被等价为tuple(torch.tensor)蒸馏数据量化算法的细节ULQ、ARQ可进一步阅读 algorithm_brief.md简易配置文件字段的权威定义以 distill_config_pytorch.proto 与 蒸馏简易配置文件 为准。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考