TVM tutorials代码学习-from_darknet.py

准备环境

要运行tvm源码下tutorials/frontend/from_darknet.py代码需要额外安装cffi和cv2包

  pip install cffi
  pip install opencv-python

详解

#导入包
import numpy as np #计算库
import matplotlib.pyplot as plot #画图
import sys
# tvm和relay IR包
import tvm
from tvm import relay
from ctypes import * #创建和执行C语言的数据类型
#darknet相关util库
from tvm.contrib.download import download_testdata
from tvm.relay.testing.darknet import __darknetffi__
import tvm.relay.testing.yolo_detection
import tvm.relay.testing.darknet

#model 名，支持'yolov2', 'yolov3' or 'yolov3-tiny'
MODEL_NAME='yolov3'
#下载cfg、weight、libdarknet2.0.so
... ...
#加载动态库导入网络
DARKNET_LIB = __darknetffi__.dlopen(lib_path)
net = DARKNET_LIB.load_network(cfg_path.encode('utf-8'), weights_path.encode('utf-8'), 0)#直接调动c库函数load网络

... ...
#1. 转换darknet到relay IR
mod, params = relay.frontend.from_darknet(net, dtype=dtype,shape=data.shape)
with relay.build_config(opt_level=3): #优化级别3
	graph, lib, params = relay.build(mod,target='llvm',target_host='llvm',params=params)
#加载测试图像
#2. 执行tvm runtime
from tvm.contrib import graph_runtime
#创建GraphModule
m = graph_runtime.create(graph,lib,tvm.cpu(0))#graph描述文件，lib文件，上下文
#设置输入data和权重
m.set_input('data', tvm.nd.array(data.astype(dtype)))
m.set_input(**params)
#执行inference
m.run()
#后处理
tvm_out = []
if MODEL_NAME = 'yolov3':
	#获取输出
	for i in range(3):
		layer_out = {}
		layer_out['type'] = 'Yolo'
		# Get the yolo layer attributes (n, out_c, out_h, out_w, classes, total)
        layer_attr = m.get_output(i*4+3).asnumpy()
        layer_out['biases'] = m.get_output(i*4+2).asnumpy()
        layer_out['mask'] = m.get_output(i*4+1).asnumpy()
        out_shape = (layer_attr[0], layer_attr[1]//layer_attr[0],
                     layer_attr[2], layer_attr[3])
        layer_out['output'] = m.get_output(i*4).asnumpy().reshape(out_shape)
        layer_out['classes'] = layer_attr[4]
        tvm_out.append(layer_out)
 #getbox
dets = tvm.relay.testing.yolo_detection.fill_network_boxes((netw, neth), (im_w, im_h), thresh,1, tvm_out)
last_layer = net.layers[net.n - 1]
#nms
tvm.relay.testing.yolo_detection.do_nms_sort(dets, last_layer.classes, nms_thresh)
#drawbox
tvm.relay.testing.yolo_detection.draw_detections(font_path, img, dets, thresh, names, last_layer.classes)
#显示
plt.imshow(img.transpose(1, 2, 0))
plt.show()

小结

整个代码流程很清晰，tvm支持大多数主流框架利用原框架的api加载model，并利用tvm前端解析器导入并编译成relay ir。

Created with Raphaël 2.2.0 darknet前端导入模型 编译成relay ir 创建tvm runtime返回GraphModule GraphModule设置输入和权值 GraphModule执行inference GraphModule获取输出feature map 后处理