CANN/AMCT量化模型接口

quantize_model

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

产品支持情况

功能说明

训练后量化接口，将输入的待量化的图结构按照给定的量化配置文件进行量化处理，在传入的图结构中插入权重量化、数据量化相关的算子，生成量化因子记录文件record_file，返回修改后的torch.nn.Module校准模型。

函数原型

calibration_model = quantize_model(config_file, modfied_onnx_file, record_file, model, input_data, input_names=None, output_names=None, dynamic_axes=None)

参数说明

返回值说明

返回修改后的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)]) scale_offset_record_file = os.path.join(TMP, 'scale_offset_record.txt') modfied_model = os.path.join(TMP, 'modfied_model.onnx') # 插入量化API calibration_model = amct.quantize_model(config_json_file, modfied_model, scale_offset_record_file, model, input_data, input_names=['input'], output_names=['output'], dynamic_axes={'input':{0: 'batch_size'}, 'output':{0: 'batch_size'}})

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