TVM tutorials代码学习-from_darknet.py

﹏ヽ暗。殇╰゛Y 2023-06-02 15:52 8阅读 0赞

### 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.0darknet前端导入模型编译成relay ir创建tvm runtime返回GraphModuleGraphModule设置输入和权值GraphModule执行inferenceGraphModule获取输出feature map后处理