环境要求

  1. python3

  2. PyTorch >= 1.8.0

    总体使用流程

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

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

    支持的op，使用建议

  5. 目前支持torch.nn.Conv2d，torch.nn.Linear的量化

  6. 仅对权值进行量化，一般8bit权值量化并不需要finetune，直接用MNN的转换工具的“—weightQuantBits”进行转换即可，但也可使用本工具进行测试精度，或者finetune到更低的bit数，如4bit，2bit，推理速度和float一致

  7. 可与剪枝工具叠加使用

    使用方法

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

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

你的模型代码

class Net(nn.Module): pass

model = Net()

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

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

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

更多配置请看API部分

quantizer = WeightQuantizer(model, bits=8) quant_model = quantizer.convert() quant_model.to(device)

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

  1. # 每次训练之前加上这一句，准备好量化训练图
  2. quantizer.resume_wq_graph()
  4. train(quant_model, data, optimizer)
  5. test(quant_model, data)
  8. if 触发模型保存条件：
  10.     # 保存模型之前去掉插入的节点，恢复原模型结构
  11.     quantizer.strip_wq_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("compress_params_index.bin", append=False)
  2. 5. 训练完之后，使用MNN转换器进行模型转换，提供上面得到的MNN压缩参数文件，并使用`--weightQuantBits`指定量化比特数：
  3. ```bash
  4. mnnconvert --modelFile quant_model_index.onnx  --MNNModel weight_quant_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params_index.bin --weightQuantBits 8

相关API

WeightQuantizer

  1. WeightQuantizer(model, bits = 8, debug_info = False, mode = 'symmetric')
  1. 参数
  3. model：Module，pytorch模型
  5. bits：int，指定权值量化比特数
  7. debug_info：bool，是否输出debug信息
  9. mode："symmetric"或"asymmetric"，采用对称或者非对称量化，目前仅支持对称量化
  1. 方法和属性
  3. convert()：返回用于训练量化的模型
  5. strip_wq_ops()：去掉插入的权值量化op，恢复原始模型结构
  7. resume_wq_graph()：strip_wq_ops保存模型之后，如果还要继续量化训练需加上这一句，以恢复量化训练图，由此支持边训练边保存
  9. save_compress_params(filename, append=False)：
  10.         用于保存MNN转换时需要用的模型压缩信息
  11.         filename：str，MNN模型压缩参数将保存到这个文件名指定的文件中
  12.         append：bool，是否将量化参数追加到filename文件中。如果进行量化的模型有剪枝，请将剪枝时通过save_compress_params生成的剪枝信息文件通过此参数传入，并将 append 设置为True