tvm 学习笔记（四）：tvm量化pytorch的onnx模型

1、模型来源：

github链接：https://github.com/qfgaohao/pytorch-ssd

2、保存成onnx模型

运行convert_to_caffe2_models.py脚本

>> python convert_to_caffe2_models.py

3、使用onnx-simplifier工具：

安装使用：

>> pip3 install onnx-simplifier
>> python3 -m onnxsim input_model output_model

4、模型转换

import onnx
import numpy as np

import tvm
import tvm.relay as relay
import tvm.relay.testing
from tvm import autotvm
from tvm.contrib import graph_runtime
import sys

target = tvm.target.create('llvm')
onnx_model = onnx.load('ssd.onnx')
input_name = '0'  # change '1' to '0'
shape_dict = {input_name: (1, 3, 300, 300)}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)

with relay.quantize.qconfig(skip_k_conv=0, round_for_shift=True):
    sym = relay.quantize.quantize(sym, params)


with relay.build_config(opt_level=1):
    graph, lib, params = relay.build_module.build(sym, target, params=params)

dtype = 'float32'
print("Output model files")
libpath = "./ssd.so"
lib.export_library(libpath)

graph_json_path = "./ssd.json"
with open(graph_json_path, 'w') as fo:
    fo.write(graph)

param_path = "./ssd.params"
with open(param_path, 'wb') as fo:
    fo.write(relay.save_param_dict(params))

5、模型部署

#!/usr/bin/python3
 
import os
import tvm
import numpy as np
import time 
from tvm.contrib.download import download
from tvm.contrib import graph_runtime
import cv2 as cv
from util.process import preprocess, non_maximum_suppression

class_names = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair",
               "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant",
               "sheep", "sofa", "train", "tvmonitor"]

test_json = 'ssd.json'
test_lib = 'ssd.so'
test_param = 'ssd.params'
 
loaded_json = open(test_json).read()
loaded_lib = tvm.module.load(test_lib)
loaded_params = bytearray(open(test_param, "rb").read())

dshape = (1, 3, 300, 300)
dtype = "float32"
ctx = tvm.cpu(0)
module = graph_runtime.create(loaded_json, loaded_lib, ctx)
module.load_params(loaded_params)

test_image_path = 'dog.jpg'
img_src = cv.imread(test_image_path)
height, width, _ = img_src.shape
img_input = preprocess(img_src)
# run the module
module.set_input("0", img_input)
module.run()
tvm_out_prob = module.get_output(0).asnumpy()
tvm_out_loc = module.get_output(1).asnumpy()

scores = tvm_out_prob[0]
boxes = tvm_out_loc[0]

picked_box_probs = []
picked_labels = []
# each time get one column 
for class_index in range(1, scores.shape[1]):
    probs = scores[:, class_index]
    mask = probs > 0.4
    probs = probs[mask]
    if probs.shape[0] == 0:
        continue
    subset_boxes = boxes[mask, :]
    # condate to a N x 5 matrix
    boxes_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], 1)
    boxes_probs = non_maximum_suppression(boxes_probs, 0.45)
    picked_box_probs.append(boxes_probs)
    picked_labels.extend([class_index] * boxes_probs.shape[0])
picked_box_probs = np.concatenate(picked_box_probs)
picked_box_probs[:, 0] *= width
picked_box_probs[:, 1] *= height
picked_box_probs[:, 2] *= width
picked_box_probs[:, 3] *= height

for box_prob, label_index in zip(picked_box_probs, picked_labels):
  cv.rectangle(img_src, (box_prob[0], box_prob[1]), (box_prob[2], box_prob[3]), (255, 0, 255), 2)
  label = "{}, {:.2f}".format(class_names[label_index], box_prob[4])
  cv.putText(img_src, label,
             (box_prob[0], box_prob[1]),
             cv.FONT_HERSHEY_SIMPLEX,
             1,
             (255, 0, 255),
             2)

path = "output.jpg"
cv.imwrite(path, img_src)
print("write image successed.")

原始效果：

量化后效果：

效果感觉有点渣，可能有什么误会，我再研究一下~

参考资料：

https://github.com/qfgaohao/pytorch-ssd