环境要求

  1. python3

  2. PyTorch >= 1.8.0

    总体使用流程

  3. 通过本工具得到训练好的float onnx模型，以及MNN模型压缩参数文件

  4. 通过MNN转换工具，输入这两个文件，得到最终的MNN量化模型

    支持的op，使用建议

  5. 建议优先试验此工具的离线量化功能，精度如果可以则不需要训练量化，速度快，节省时间

  6. 目前支持torch.nn.Conv2d，torch.nn.Linear的量化
  7. 优化器超参使用模型收敛阶段的超参，学习率可以在收敛阶段学习率的基础上再调小
  8. 模型的第一层或者前面几层对精度影响较大，可以尝试跳过不进行量化

  9. relu6建议使用nn.ReLU6，F.relu6使用opset 9导出到onnx会有问题

    支持的训练量化算法

  10. LSQQuantizer: 改进naive QAT算法，scale会在网络训练中学习更新

    使用方法

  11. 建议从训练好的float模型进行finetune

  12. 创建Quantizer对象，对原始模型进行转换，然后用转换之后的模型进行训练
  13. 保存onnx模型之前，去掉插入的节点
  14. 保存onnx模型之后，导出MNN模型压缩参数文件，示例代码如下（关注其中quantizer的用法）： ```python from mnncompress.pytorch import LSQQuantizer Quantizer = LSQQuantizer

你的模型代码

class Net(nn.Module): pass

model = Net()

加载已经训练好的模型，可以是剪枝之后的

model.load_state_dict(torch.load(“pruned_model.pt”))

将模型进行转换，并使用转换后的模型进行训练，测试

retain_sparsity=True表示待量化的float模型是稀疏模型，希望叠加量化训练

更多配置请看API部分

注意在model还在cpu上，任何分布式还未生效前调用

quantizer = Quantizer(model, retain_sparsity=False) quant_model = quantizer.convert()

单机或分布式，根据你已有的代码来

quant_model.to(device)

for epoch in range(1, epochs + 1):

  1. # 每次训练之前加上这一句，准备好量化训练图
  2. quantizer.resume_qat_graph()
  4. train(quant_model, data, optimizer)
  5. test(quant_model, data)
  8. if 触发模型保存条件：
  10.     # 保存模型之前去掉插入的节点，恢复原模型结构
  11.     quantizer.strip_qat_ops()
  13.     # 保存模型，注意index，即模型和保存MNN模型压缩参数文件是一一对应的
  14.     quant_model.eval()
  15.     torch.save(quant_model.state_dict(), "quant_model_index.pt")
  16.     x = torch.randn(input_shape).to(device)
  17.     torch.onnx.export(quant_model, x, "quant_model_index.onnx")
  19.     # 保存MNN模型压缩参数文件，如果进行量化的模型有剪枝，
  20.     # 请将剪枝时生成的MNN模型压缩参数文件 "compress_params.bin" 文件在下方传入，并将 append 设置为True
  21.     quantizer.save_compress_params("quant_model_index.onnx", "compress_params_index.bin", append=False)
  2. 5. 将onnx模型和生成的 "compress_params.bin" 文件输入到MNN转换工具中进行转换，得到最终的MNN量化模型：
  3. ```bash
  4. mnnconvert --modelFile quant_model.onnx  --MNNModel quant_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params.bin

作为离线量化工具使用

去掉训练过程中的参数更新部分，并将LSQQuantizer的mode设置为offline（查看下方API），其他步骤和上述一致，此时该工具即成为离线量化工具，直接灌入训练数据即可完成离线量化：

  1. # 注释掉训练部分中的这一句即可
  2. # optimizer.step()

可以先尝试这个，如果精度可接受，可以节省训练时间。

相关API

LSQQuantizer

  1. LSQQuantizer(model, skip_quant_layers = [], bits = 8, debug_info = False, mode = 'online', retain_sparsity=False)
  1. 参数
  3. model：Module，pytorch模型
  5. skip_quant_layers：[str,]，不进行量化的module的名字，可通过dict(model.named_modules()).keys()获得，如['model.conv1', 'model.conv2']，需要是nn.Conv2d或者nn.Linear
  7. bits：int，指定量化比特数
  9. debug_info：bool，是否输出debug信息
  11. mode："online"或"offline"，训练量化时选online，离线量化时选offline
  13. retain_sparsity: bool，为True表示待量化的float模型是稀疏模型，希望叠加量化训练
  1. 方法和属性
  3. convert()：返回用于训练量化的模型
  5. strip_qat_ops()：去掉插入的训练量化op，恢复原始模型结构
  7. resume_qat_graph()：strip_qat_ops保存模型之后，如果还要继续量化训练需加上这一句，以恢复量化训练图，由此支持边训练边保存
  9. save_compress_params(onnx_inference_model_file, filename, append=False)：
  10.         用于保存MNN转换时需要用的模型压缩信息
  11.         onnx_inference_model_file：str，去掉插入的训练量化节点之后保存的onnx推理模型名字
  12.         filename：str，MNN模型压缩参数将保存到这个文件名指定的文件中
  13.         append：bool，是否将量化参数追加到filename文件中。如果进行量化的模型有剪枝，请将剪枝时通过save_compress_params生成的剪枝信息文件通过此参数传入，并将 append 设置为True