训练后量化工具需要用户提供PyTorch训练脚本或者pth文件,工具可自动对模型中的卷积(torch.nn.Conv2d)和线性层(torch.nn.Linear)进行识别并量化,最终导出量化后的onnx模型,量化后的模型可以在推理服务器上运行,达到提升推理性能的目的。量化过程中用户需自行提供模型与数据集,调用API接口完成模型的量化调优。
为了提升量化精度,训练后量化(PyTorch)算法内置了自动混合精度的模块,自动识别并回退量化敏感层为浮点计算,避免量化敏感层对精度造成较大损失。算法核心是:计算每个量化层量化前后输出的MSE,根据MSE的排序来衡量每一个量化层的量化敏感性,自动回退MSE最大的部分敏感层,从而提升量化的精度。
为了进一步降低量化精度损失,训练后量化(PyTorch)工具内集成了多种精度保持策略,对权重的量化参数和取证方式进行优化。
如下命令如果使用非root用户安装,需要在安装命令后加上--user,例如:pip3 install onnx --user。
pip3 install numpy #需大于等于1.21.0版本 pip3 install onnx #需大于等于1.11.0版本 pip3 install torch==1.11.0 #支持1.8.1和1.11.0,须为CPU版本的torch pip3 install onnx-simplifier #需大于等于0.3.10版本
量化脚本中需用户关注以下步骤,可以根据实际情况配置。
import torchvision model = torchvision.models.resnet50(pretrained=True) model.eval()
from msmodelslim.pytorch.quant.ptq_tools import QuantConfig, Calibrator
from msmodelslim import set_logger_level set_logger_level("info") #根据实际情况配置
disable_names = [] input_shape = [1, 3, 224, 224] keep_acc={'admm': [False, 1000], 'easy_quant': [False, 1000], 'round_opt': False} quant_config = QuantConfig( disable_names=disable_names, # 手动回退的量化层名称,如精度太差,推荐回退量化敏感层 amp_num=0, # 混合精度量化回退层数 input_shape=input_shape, # 模型输入的shape,用于Label-Free量化构造虚拟数据 keep_acc=keep_acc, # 精度保持策略 sigma=25, # 大于0使用sigma统计方法;传入0值使用min-max统计方法 )
calibrator = Calibrator(model, quant_config)
calibrator.run()
calibrator.export_param ("./output")
calibrator.export_quant_onnx("resnet50", "./output", ["input.1"])
python3 resnet_quant.py