CANN/amct蒸馏模型保存接口
2026/5/10 6:14:17 网站建设 项目流程

save_distill_model

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

产品支持情况

产品是否支持
Ascend 950PR/Ascend 950DT
Atlas A3 训练系列产品/Atlas A3 推理系列产品
Atlas A2 训练系列产品/Atlas A2 推理系列产品

功能说明

蒸馏接口,根据用户最终的蒸馏好的模型,生成最终量化精度仿真模型以及量化部署模型。

函数原型

save_distill_model(model, save_path, input_data, record_file=None, input_names=None, output_names=None, dynamic_axes=None)

参数说明

参数名

输入/输出

说明

model

输入

含义:已进行蒸馏后的量化模型。

数据类型:torch.nn.Module

save_path

输入

含义:蒸馏后的量化模型存放路径。该路径需要包含模型名前缀,例如./quantized_model/*model。

数据类型:string

input_data

输入

含义:模型的输入数据。一个torch.tensor会被等价为tuple(torch.tensor)。

数据类型:tuple

record_file

输入

含义:量化因子记录文件路径及名称。

默认值:None

数据类型:string

使用约束:传入值为None的情况下量化因子记录文件存放在amct_log文件夹中。

input_names

输入

含义:模型输入节点的名称,用于在保存的量化onnx模型中显示。

默认值:None

数据类型:list(string)

output_names

输入

含义:模型输出节点的名称,用于在保存的量化onnx模型中显示。

默认值:None

数据类型:list(string)

dynamic_axes

输入

含义:对模型输入输出动态轴的指定,例如对于输入inputs(NCHW),N、H、W为不确定大小,输出outputs(NL),N为不确定大小,则指定形式为:{"inputs": [0,2,3], "outputs": [0]},其中0,2,3分别表示N,H,W所在位置的索引。

默认值:None

数据类型:dict<string, dict<python:int, string>> or dict<string, list(int)>

返回值说明

调用示例

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,将蒸馏的模型存为onnx文件 amct.save_distill_model( model, "./model/distilled", input_data, record_file="./results/records.txt", input_names=['input'], output_names=['output'], dynamic_axes={'input':{0: 'batch_size'}, 'output':{0: 'batch_size'}})

落盘文件说明:

  • 精度仿真模型文件:ONNX格式的模型文件,模型名中包含fake_quant,可以在ONNX Runtime环境进行精度仿真。

  • 部署模型文件:ONNX格式的模型文件,模型名中包含deploy,经过ATC转换工具转换后可部署到AI处理器。

重新执行蒸馏时,该接口输出的上述文件将会被覆盖。

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

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询