tvm学习笔记（二）：载入insightface人脸识别模型

1、载入insightface训练好的人脸识别模型

将insightface训练好的模型转换成tvm能部署的model格式：

import numpy as np
import nnvm.compiler
import nnvm.testing
import tvm
from tvm.contrib import graph_runtime
import mxnet as mx
from mxnet import ndarray as nd

prefix, epoch = "model", 0
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
image_size = (112, 112)
opt_level = 3

shape_dict = {'data':(1, 3, *image_size)}
target = tvm.target.create("llvm -mcpu=haswell")
# target = tvm.target.create("llvm -mcpu=broadwell")

nnvm_sym, nnvm_params = nnvm.frontend.from_mxnet(sym, arg_params, aux_params)
with nnvm.compiler.build_config(opt_level=opt_level):
    graph, lib, params = nnvm.compiler.build(nnvm_sym, target, shape_dict, params=nnvm_params)

lib.export_library("./deploy_lib.so")
print('lib export successfully')
with open("./deploy_graph.json", "w") as fo:
    fo.write(graph.json())
with open("./deploy_param.params", "wb") as fo:
    fo.write(nnvm.compiler.save_param_dict(params))

运行之后，可以生成deploy_lib.so，deploy_graph.json， deploy_param.params三个模型文件。

2、对模型进行部署

import numpy as np
import nnvm.compiler
import nnvm.testing
import tvm
from tvm.contrib import graph_runtime
import mxnet as mx
from mxnet import ndarray as nd

dtype = "float32"
ctx = tvm.cpu()
loaded_json = open("./deploy_graph.json").read()
loaded_lib = tvm.module.load("./deploy_lib.so")
loaded_params = bytearray(open("./deploy_param.params", "rb").read())

data_shape = (1, 3, 112, 112)
import cv2
face_1 = cv2.imread("face1.png")
face_2 = cv2.imread("face2.png")
print(data_shape)
face_1 = cv2.resize(face_1, (data_shape[2], data_shape[3]))
face_2 = cv2.resize(face_2, (data_shape[2], data_shape[3]))
face_1 = np.transpose(np.array(face_1), (2, 0, 1))
face_2 = np.transpose(np.array(face_2), (2, 0, 1))

input_face1 = tvm.nd.array(face_1.astype(dtype))
input_face2 = tvm.nd.array(face_2.astype(dtype))
module = graph_runtime.create(loaded_json, loaded_lib, ctx)
module.load_params(loaded_params)


module.run(data=input_face1)
v1 = module.get_output(0).asnumpy()[0]
module.run(data=input_face2)
v2 = module.get_output(0).asnumpy()[0]
num=float(np.sum(v1*v2))
denom=np.linalg.norm(v1)*np.linalg.norm(v2)
cos = num / denom

print("***************")
print(v1)
print("***************")
print(v2)
print("***************")
print("dist: ", cos)

Everything:

https://github.com/MirrorYuChen/tvm_insightface.git

参考资料：

https://github.com/deepinsight/insightface/wiki/Tutorial:-Deploy-Face-Recognition-Model-via-TVM