AI編譯器TVM部署示例解析

AI編譯器TVM部署示例解析

AI編譯器TVM(一)——一個簡單的例子

概述

什么是TVM？

TVM可以稱為許多工具集的集合，這些工具可以組合起來使用，實現一些神經網絡的加速和部署功能。這也是為什么叫做TVM Stack了。TVM的使用途徑很廣，幾乎可以支持市面上大部分的神經網絡權重框架(ONNX、TF、Caffe2等)，也幾乎可以部署在任何的平台，如Windows、Linux、Mac、ARM等等。

參考文獻

https://oldpan.me/archives/the-first-step-towards-tvm-1

https://mp.weixin.qq.com/s?__biz=Mzg3ODU2MzY5MA==&mid=2247484929&idx=1&sn=3fcce36b5a50cd8571cf932a23083667&chksm=cf109e04f86717129c3381ebeec2d0c1f7baf6ed057c66310662f5935beea88baf23e99898f4&token=1276531538&lang=zh_CN#rd

https://mp.weixin.qq.com/s?__biz=Mzg3ODU2MzY5MA==&mid=2247484930&idx=1&sn=ddc3da7b72c900ce2f8e6aad99a9e788&source=41#wechat_redirect

以下面一張圖來形容一下，這張圖來源於(https://tvm.ai/about)：

 

 stack_tvmlang

只需要知道TVM的核心功能就可以：TVM可以優化的訓練好的模型，將模型打包好，將這個優化好的模型放在任何平台去運行，可以說是與落地應用息息相關。

TVM包含的東西和知識概念都有很多，不僅有神經網絡優化量化op融合等一系列步驟，還有其他更多細節技術的支持(Halide、LLVM)，從而使TVM擁有很強大的功能。如果想多了解TVM的可以在知乎上直接搜索TVM關鍵字，那些大佬有很多關於TVM的介紹文章，大家可以去看看。

其實做模型優化這一步驟的庫已經出現很多了，不論是Nvidia自家的TensorRT，還是Pytorch自家的torch.jit模塊，都在做一些模型優化的工作，這里就不多說了，感興趣的可以看看以下文章：

利用Pytorch的C++前端(libtorch)讀取預訓練權重並進行預測
利用TensorRT實現神經網絡提速(讀取ONNX模型並運行)
利用TensorRT對深度學習進行加速

開始使用

為什么要使用TVM？

如果想將訓練模型移植到Window端、ARM端(樹莓派、其他一系列使用該內核的板卡)或者其他的一些平台，利用其中的CPU或者GPU來運行，希望可以通過優化模型來使模型，在該平台運算的速度更快(這里與模型本身的算法設計無關)，實現落地應用研究，那么TVM就是不二之選。另外TVM源碼是由C++和Pythoh共同搭建，閱讀相關源碼也有利於程序編寫方面的提升。

安裝

安裝其實沒什么多說的，官方的例子說明的很詳細。大家移步到那里按照官方的步驟一步一步來即可。

不過有兩點需要注意下：

• 建議安裝LLVM，雖然LLVM對於TVM是可選項，但是如果想要部署到CPU端，那么llvm幾乎是必須的
• 因為TVM是python和C++一起的工程，python可以說是C++的前端，安裝官方教程編譯好C++端后，這里建議選擇官方中的Method 1來進行python端的設置，這樣就可以隨意修改源代碼，再重新編譯，Python端就不需要進行任何修改就可以直接使用了。

(官方建議使用Method 1)

利用Pytorch導出Onnx模型

這里以一個簡單的例子，演示一下TVM是怎么使用的。

首先要做的是，得到一個已經訓練好的模型，這里選擇這個github倉庫中的mobilenet-v2，model代碼和在ImageNet上訓練好的權重都已經提供。將github中的模型代碼移植到本地，然后調用並加載已經訓練好的權重：

import torch

import time

from models.MobileNetv2 import mobilenetv2  


model = mobilenetv2(pretrained=True)

example = torch.rand(1, 3, 224, 224)   # 假想輸入


with torch.no_grad():

    model.eval()

    since = time.time()

    for i in range(10000):

        model(example)

    time_elapsed = time.time() - since

    print('Time elapsed is {:.0f}m {:.0f}s'.

          format(time_elapsed // 60, time_elapsed % 60))  # 打印出來時間

這里加載訓練好的模型權重，設定了輸入，在python端連續運行了10000次，這里所花的時間為：6m2s。

然后將Pytorch模型導出為ONNX模型：import torch

from models.MobileNetv2 import mobilenetv2  


model = mobilenetv2(pretrained=True)

example = torch.rand(1, 3, 224, 224)   # 假想輸入


torch_out = torch.onnx.export(model,

                              example,

                              "mobilenetv2.onnx",

                              verbose=True,

                              export_params=True   # 帶參數輸出

                              )

這樣就得到了mobilenetv2.onnx這個onnx格式的模型權重。這里要帶參數輸出，因為之后要直接讀取ONNX模型進行預測。

導出來之后，建議使用Netron來查看模型的結構，可以看到這個模型由Pytorch-1.0.1導出，共有152個op，以及輸入id和輸入格式等等信息，可以拖動鼠標查看到更詳細的信息:

 

 mobilenetv2-test

至此mobilenet-v2模型已經順利導出了。

利用TVM讀取並預測ONNX模型

在成功編譯並且可以在Python端正常引用TVM后，首先導入onnx格式的模型。這里准備了一張飛機的圖像：

tvm_plane

 

 這個圖像在ImageNet分類中屬於404: 'airliner'，也就是航空客機。

下面將利用TVM部署onnx模型並對這張圖像進行預測。import onnx

import time

import tvm

import numpy as np

import tvm.relay as relay

from PIL import Image


onnx_model = onnx.load('mobilenetv2.onnx')  # 導入模型


mean = [123., 117., 104.]                   # 在ImageNet上訓練數據集的mean和std

std = [58.395, 57.12, 57.375]



def transform_image(image):                # 定義轉化函數，將PIL格式的圖像轉化為格式維度的numpy格式數組

    image = image - np.array(mean)

    image /= np.array(std)

    image = np.array(image).transpose((2, 0, 1))

    image = image[np.newaxis, :].astype('float32')

    return image


img = Image.open('../datasets/images/plane.jpg').resize((224, 224)) # 這里將圖像resize為特定大小

x = transform_image(img)

這樣得到的x為[1,3,224,224]維度的ndarray。這個符合NCHW格式標准，也是通用的張量格式。

接下來設置目標端口llvm，也就是部署到CPU端，這里使用的是TVM中的Relay IR，這個IR簡單來說就是可以讀取模型，按照模型的順序搭建出一個可以執行的計算圖，可以對這個計算圖進行一系列優化。(現在TVM主推Relay而不是NNVM，Relay可以稱為二代NNVM)。

target = 'llvm'


input_name = '0'  # 注意這里為之前導出onnx模型中的模型的輸入id，這里為0

shape_dict = {input_name: x.shape}

# 利用Relay中的onnx前端讀取導出的onnx模型

sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)

上述代碼中導出的sym和params是接下來要使用的核心的東西，其中params就是導出模型中的權重信息，在python中用dic表示：

Screenshot from 2019-03-12 14-57-18

sym就是表示計算圖結構的功能函數，這個函數中包含了計算圖的流動過程，以及一些計算中需要的各種參數信息，Relay IR之后對網絡進行優化就是主要對這個sym進行優化的過程：

fn (%v0: Tensor[(1, 3, 224, 224), float32],

    %v1: Tensor[(32, 3, 3, 3), float32],

    %v2: Tensor[(32,), float32],

    %v3: Tensor[(32,), float32],

    %v4: Tensor[(32,), float32],

    %v5: Tensor[(32,), float32],

    ...

    %v307: Tensor[(1280, 320, 1, 1), float32],

    %v308: Tensor[(1280,), float32],

    %v309: Tensor[(1280,), float32],

    %v310: Tensor[(1280,), float32],

    %v311: Tensor[(1280,), float32],

    %v313: Tensor[(1000, 1280), float32],

    %v314: Tensor[(1000,), float32]) {

  %0 = nn.conv2d(%v0, %v1, strides=[2, 2], padding=[1, 1], kernel_size=[3, 3])

  %1 = nn.batch_norm(%0, %v2, %v3, %v4, %v5, epsilon=1e-05)

  %2 = %1.0

  %3 = clip(%2, a_min=0, a_max=6)

  %4 = nn.conv2d(%3, %v7, padding=[1, 1], groups=32, kernel_size=[3, 3])

  ...

  %200 = clip(%199, a_min=0, a_max=6)

  %201 = mean(%200, axis=[3])

  %202 = mean(%201, axis=[2])

  %203 = nn.batch_flatten(%202)

  %204 = multiply(1f, %203)

  %205 = nn.dense(%204, %v313, units=1000)

  %206 = multiply(1f, %v314)

  %207 = nn.bias_add(%205, %206)

  %207

}

接下來需要對這個計算圖模型進行優化，這里選擇優化的等級為3：

with relay.build_config(opt_level=3):

    intrp = relay.build_module.create_executor('graph', sym, tvm.cpu(0), target)


dtype = 'float32'

func = intrp.evaluate(sym)

最后，得到可以直接運行的func。

其中優化的等級分這幾種：

OPT_PASS_LEVEL = {

    "SimplifyInference": 0,

    "OpFusion": 1,

    "FoldConstant": 2,

    "CombineParallelConv2D": 3,

    "FoldScaleAxis": 3,

    "AlterOpLayout": 3,

    "CanonicalizeOps": 3,

}

最后，將之前已經轉化格式后的圖像x數組和模型的參數輸入到這個func中，返回這個輸出數組中的最大值

output = func(tvm.nd.array(x.astype(dtype)), **params).asnumpy()

print(output.argmax())

這里得到的輸出為404，與前文描述圖像在ImageNet中的分類標記一致，說明TVM正確讀取onnx模型並將其應用於預測階段。

另外單獨測試一下模型優化后運行的速度和之前直接利用pytorch運行速度之間比較一下，最后的運行時間為：3m20s，相較之前的6m2s快了將近一倍。

since = time.time()

for i in range(10000):

    output = func(tvm.nd.array(x.astype(dtype)), **params).asnumpy()

time_elapsed = time.time() - since

print('Time elapsed is {:.0f}m {:.0f}s'.

      format(time_elapsed // 60, time_elapsed % 60))  # 打印出來時間

當然，這個比較並不是很規范，不過可以大概分析出TVM的一些可用之處了。

這里了解一下什么是TVM以及一個簡單例子的使用，在接下來會涉及到部分TVM設計結構和源碼的解析。可能涉及到的知識點有：

• 簡單編譯器原理
• C++特殊語法以及模板元編程
• 神經網絡模型優化過程
• 代碼部署

等等，隨時可能會進行變化。

人工智能已經開始進入嵌入式時代，各式各樣的AI芯片即將初始，將復雜的網絡模型運行在廉價低功耗的板子上可能也不再是遙不可及的幻想，不知道未來會是怎么樣，但TVM這個框架已經開始走了一小步。

AI編譯器TVM(二)——利用TVM完成C++端的部署

前言

在上一節，簡單介紹了什么是TVM以及如何利用Relay IR去編譯網絡權重然后並運行起來。

 

 TVM

上述文章中的例子很簡單，但是實際中更需要的是利用TVM去部署應用么，最簡單直接的就是在嵌入式系統中運行起神經網絡模型。例如樹莓派。這才是最重要的是不是？所以嘛，在深入TVM之前還是要走一遍基本的實踐流程的，也唯有實踐流程才能讓更好地理解TVM到底可以做什么。

本節主要介紹如果將神經網絡使用TVM編譯，導出動態鏈接庫文件，最后部署在樹莓派端(PC端)，運行起來。

環境搭建

環境搭建？有什么好講的？

廢話咯，需要先把TVM的環境搭建出來才可以用啊，官方的安裝教程最為詳細，這里還是多建議看看官方的文檔，很詳細很具體重點把握的也很好。

但是還是要強調兩點：

• 需要安裝LLVM，因為這篇文章所講的主要運行環境是CPU(樹莓派的GPU暫時不用，內存有點小)，所以LLVM是必須的
• 安裝交叉編譯器:

Cross Compiler

交叉編譯器是什么，就是可以在PC平台上編譯生成可以直接在樹莓派上運行的可執行文件。在TVM中，需要利用交叉編譯器在PC端編譯模型並且優化，然后生成適用於樹莓派(arm構架)使用的動態鏈接庫。

有這個動態鏈接庫，就可以直接調用樹莓派端的TVM運行時環境去調用這個動態鏈接庫，執行神經網絡的前向操作了。

怎么安裝呢？這里需要安裝叫做/usr/bin/arm-linux-gnueabihf-g++的交叉編譯器，在Ubuntu系統中，直接sudo apt-get install g++-arm-linux-gnueabihf即可，注意名稱不能錯，需要的是hf(Hard-float)版本。

安裝完后，執行/usr/bin/arm-linux-gnueabihf-g++ -v命令就可以看到輸出信息:

1prototype@prototype-X299-UD4-Pro:~/$ /usr/bin/arm-linux-gnueabihf-g++ -v

2Using built-in specs.

3COLLECT_GCC=/usr/bin/arm-linux-gnueabihf-g++

4COLLECT_LTO_WRAPPER=/usr/lib/gcc-cross/arm-linux-gnueabihf/5/lto-wrapper

5Target: arm-linux-gnueabihf

6Configured with: ../src/configure -v --with-pkgversion='Ubuntu/Linaro 5.4.0-6ubuntu1~16.04.9' --with-bugurl=file:///usr/share/doc/gcc-5/README.Bugs --enable-languages=c,ada,c++,java,go,d,fortran,objc,obj-c++ --prefix=/usr --program-suffix=-5 --enable-shared --enable-linker-build-id --libexecdir=/usr/lib --without-included-gettext --enable-threads=posix --libdir=/usr/lib --enable-nls --with-sysroot=/ --enable-clocale=gnu --enable-libstdcxx-debug --enable-libstdcxx-time=yes --with-default-libstdcxx-abi=new --enable-gnu-unique-object --disable-libitm --disable-libquadmath --enable-plugin --with-system-zlib --disable-browser-plugin --enable-java-awt=gtk --enable-gtk-cairo --with-java-home=/usr/lib/jvm/java-1.5.0-gcj-5-armhf-cross/jre --enable-java-home --with-jvm-root-dir=/usr/lib/jvm/java-1.5.0-gcj-5-armhf-cross --with-jvm-jar-dir=/usr/lib/jvm-exports/java-1.5.0-gcj-5-armhf-cross --with-arch-directory=arm --with-ecj-jar=/usr/share/java/eclipse-ecj.jar --disable-libgcj --enable-objc-gc --enable-multiarch --enable-multilib --disable-sjlj-exceptions --with-arch=armv7-a --with-fpu=vfpv3-d16 --with-float=hard --with-mode=thumb --disable-werror --enable-multilib --enable-checking=release --build=x86_64-linux-gnu --host=x86_64-linux-gnu --target=arm-linux-gnueabihf --program-prefix=arm-linux-gnueabihf- --includedir=/usr/arm-linux-gnueabihf/include

7Thread model: posix

8gcc version 5.4.0 20160609 (Ubuntu/Linaro 5.4.0-6ubuntu1~16.04.9) 

樹莓派環境搭建

因為是在PC端利用TVM編譯神經網絡的，所以在樹莓派端只需要編譯TVM的運行時環境即可(TVM可以分為兩個部分，一部分為編譯時，另一個為運行時，兩者可以拆開)。

這里附上官方的命令，注意樹莓派端需要安裝llvm，樹莓派端的llvm可以在llvm官方找到已經編譯好的壓縮包，解壓后添加環境變量即可：

1git clone --recursive https://github.com/dmlc/tvm

2cd tvm

3mkdir build

4cp cmake/config.cmake build   # 這里修改config.cmake使其支持llvm

5cd build

6cmake ..

7make runtime

在樹莓派上編譯TVM的運行時並不需要花很久的時間。

完成部署

環境搭建好之后，就讓開始部署任務。

首先依然需要一個自己的測試模型，在這里使用之前訓練好的，識別剪刀石頭布手勢的模型權重，然后利用Pytorch導出ONNX模型出來。具體的導出步驟可以看下面這兩篇文章。

• 利用Pytorch的C++前端(libtorch)讀取預訓練權重並進行預測
• Pytorch的C++端(libtorch)在Windows中的使用

 

 

 (上圖是之前的識別剪刀石頭布的一個權重模型)

 

OK，那擁有了一個模型叫做mobilenetv2-128_S.onnx，這個模型也就是通過Pytorch導出的ONNX模型，利用Netron瞧一眼：

 

 

 整個模型的輸入和輸出上圖寫的都很清楚了。

測試模型

拿到模型后，首先測試模型是否可以正確工作，同上一篇介紹TVM的文章類似，利用TVM的PYTHON前端去讀取.onnx模型，然后將其編譯並運行，最后利用測試圖像測試其是否可以正確工作，其中核心代碼如下：

 1onnx_model = onnx.load('../test/new-mobilenetv2-128_S.onnx')

 2

 3img = Image.open('../datasets/hand-image/paper.jpg').resize((128, 128))

 4

 5img = np.array(img).transpose((2, 0, 1)).astype('float32')  

 6img = img/255.0           # 注意在Pytorch中的tensor范圍是0-1

 7x = img[np.newaxis, :]

 8

 9target = 'llvm'

10

11input_name = '0'      # 這里需要注意，因為生成的.onnx模型的輸入代號是0，所以這里改為0

12shape_dict = {input_name: x.shape}

13sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)

14

15with relay.build_config(opt_level=3):

16    intrp = relay.build_module.create_executor('graph', sym, tvm.cpu(0), target)

17

18dtype = 'float32'

19func = intrp.evaluate(sym)

20

21# 輸出推斷的結果

22tvm_output = intrp.evaluate(sym)(tvm.nd.array(x.astype(dtype)), **params).asnumpy()

23max_index = tvm_output.argmax()

24print(max_index)

這個模型輸出的結果為三個手勢的輸出值大小(順序分別為布、剪刀、石頭)，上述的代碼打印出來的值為0，意味着可以正確識別paper.jpg輸入的圖像。說明這個轉化過程是沒有問題的。

導出動態鏈接庫

上面這個步驟只是將.onnx模型利用TVM讀取並且預測出來，如果需要部署的話就需要導出整個模型的動態鏈接庫，至於為什么是動態鏈接庫，其實TVM是有多種的導出模式的(也可以導出靜態庫)，但是這里不細說了：

 

總之目標就是導出so動態鏈接庫，這個鏈接庫中包括了神經網絡所需要的一切推斷功能。

怎么導出呢？其實官方已經有很詳細的導出說明。這里不進行贅述了，僅僅展示核心的代碼加以注釋即可。

請看以下的代碼：

 1#開始同樣是讀取.onnx模型

 2

 3onnx_model = onnx.load('../../test/new-mobilenetv2-128_S.onnx')

 4img = Image.open('../../datasets/hand-image/paper.jpg').resize((128, 128))

 5

 6# 以下的圖片讀取僅僅是為了測試

 7img = np.array(img).transpose((2, 0, 1)).astype('float32')

 8img = img/255.0    # remember pytorch tensor is 0-1

 9x = img[np.newaxis, :]

10

11# 這里首先在PC的CPU上進行測試 所以使用LLVM進行導出

12target = tvm.target.create('llvm')

13

14input_name = '0'  # change '1' to '0'

15shape_dict = {input_name: x.shape}

16sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)

17

18# 這里利用TVM構建出優化后模型的信息

19with relay.build_config(opt_level=2):

20    graph, lib, params = relay.build_module.build(sym, target, params=params)

21

22dtype = 'float32'

23

24from tvm.contrib import graph_runtime

25

26# 下面的函數導出需要的動態鏈接庫 地址可以自己定義

27print("Output model files")

28libpath = "../tvm_output_lib/mobilenet.so"

29lib.export_library(libpath)

30

31# 下面的函數導出神經網絡的結構，使用json文件保存

32graph_json_path = "../tvm_output_lib/mobilenet.json"

33with open(graph_json_path, 'w') as fo:

34    fo.write(graph)

35

36# 下面的函數中導出神經網絡模型的權重參數

37param_path = "../tvm_output_lib/mobilenet.params"

38with open(param_path, 'wb') as fo:

39    fo.write(relay.save_param_dict(params))

40# -------------至此導出模型階段已經結束--------

41

42# 接下來加載導出的模型去測試導出的模型是否可以正常工作

43loaded_json = open(graph_json_path).read()

44loaded_lib = tvm.module.load(libpath)

45loaded_params = bytearray(open(param_path, "rb").read())

46

47# 這里執行的平台為CPU

48ctx = tvm.cpu()

49

50module = graph_runtime.create(loaded_json, loaded_lib, ctx)

51module.load_params(loaded_params)

52module.set_input("0", x)

53module.run()

54out_deploy = module.get_output(0).asnumpy()

55

56print(out_deploy)

上述的代碼輸出[[13.680096 -7.218611 -6.7872353]]，因為輸入的圖像是paper.jpg,所以輸出的三個數字第一個數字最大，沒有毛病。

執行完代碼之后就可以得到需要的三個文件

• mobilenet.so
• mobilenet.json
• mobilenet.params

得到三個文件之后，接下來利用TVM的C++端讀取並運行起來。

在PC端利用TVM部署C++模型

如何利用TVM的C++端去部署，官方也有比較詳細的文檔，這里利用TVM和OpenCV讀取一張圖片，並且使用之前導出的動態鏈接庫去運行神經網絡對這張圖片進行推斷。

需要的頭文件為：

1#include <cstdio>

2#include <dlpack/dlpack.h>

3#include <opencv4/opencv2/opencv.hpp>

4#include <tvm/runtime/module.h>

5#include <tvm/runtime/registry.h>

6#include <tvm/runtime/packed_func.h>

7#include <fstream>

其實這里只需要TVM的運行時，另外dlpack是存放張量的一個結構。其中OpenCV用於讀取圖片， fstream則用於讀取json和參數信息：

 1tvm::runtime::Module mod_dylib =

 2    tvm::runtime::Module::LoadFromFile("../files/mobilenet.so");

 3

 4std::ifstream json_in("../files/mobilenet.json", std::ios::in);

 5std::string json_data((std::istreambuf_iterator<char>(json_in)), std::istreambuf_iterator<char>());

 6json_in.close();

 7

 8// parameters in binary

 9std::ifstream params_in("../files/mobilenet.params", std::ios::binary);

10std::string params_data((std::istreambuf_iterator<char>(params_in)), std::istreambuf_iterator<char>());

11params_in.close();

12

13TVMByteArray params_arr;

14params_arr.data = params_data.c_str();

15params_arr.size = params_data.length();

在讀取完信息之后，要利用之前讀取的信息，構建TVM中的運行圖(Graph_runtime)：

1int dtype_code = kDLFloat;

2int dtype_bits = 32;

3int dtype_lanes = 1;

4int device_type = kDLCPU;

5int device_id = 0;

6

7tvm::runtime::Module mod = (*tvm::runtime::Registry::Get("tvm.graph_runtime.create"))

8        (json_data, mod_dylib, device_type, device_id);

然后利用TVM中函數建立一個輸入的張量類型並且分配空間：

1DLTensor *x;

2int in_ndim = 4;

3int64_t in_shape[4] = {1, 3, 128, 128};

4TVMArrayAlloc(in_shape, in_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &x);

其中DLTensor是個靈活的結構，可以包容各種類型的張量，而在創建了這個張量后，需要將OpenCV中讀取的圖像信息傳入到這個張量結構中：

1// 這里依然讀取了papar.png這張圖

2image = cv::imread("/home/prototype/CLionProjects/tvm-cpp/data/paper.png");

3

4cv::cvtColor(image, frame, cv::COLOR_BGR2RGB);

5cv::resize(frame, input,  cv::Size(128,128));

6

7float data[128 * 128 * 3];

8// 在這個函數中 將OpenCV中的圖像數據轉化為CHW的形式 

9Mat_to_CHW(data, input);

需要注意的是，因為OpenCV中的圖像數據的保存順序是(128,128,3)，所以這里需要將其調整過來，其中Mat_to_CHW函數的具體內容是:

 1void Mat_to_CHW(float *data, cv::Mat &frame)

 2{

 3    assert(data && !frame.empty());

 4    unsigned int volChl = 128 * 128;

 5

 6    for(int c = 0; c < 3; ++c)

 7    {

 8        for (unsigned j = 0; j < volChl; ++j)

 9            data[c*volChl + j] = static_cast<float>(float(frame.data[j * 3 + c]) / 255.0);

10    }

11

12}

當然別忘了除以255.0因為在Pytorch中所有的權重信息的范圍都是0-1。

在將OpenCV中的圖像數據轉化后，將轉化后的圖像數據拷貝到之前的張量類型中:

1// x為之前的張量類型 data為之前開辟的浮點型空間

2memcpy(x->data, &data, 3 * 128 * 128 * sizeof(float));

然后設置運行圖的輸入(x)和輸出(y):

 1// get the function from the module(set input data)

 2tvm::runtime::PackedFunc set_input = mod.GetFunction("set_input");

 3set_input("0", x);

 4

 5// get the function from the module(load patameters)

 6tvm::runtime::PackedFunc load_params = mod.GetFunction("load_params");

 7load_params(params_arr);

 8

 9DLTensor* y;

10int out_ndim = 2;

11int64_t out_shape[2] = {1, 3,};

12TVMArrayAlloc(out_shape, out_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &y);

13

14// get the function from the module(run it)

15tvm::runtime::PackedFunc run = mod.GetFunction("run");

16

17// get the function from the module(get output data)

18tvm::runtime::PackedFunc get_output = mod.GetFunction("get_output");

可以運行了：

1run();

2get_output(0, y);

3

4// 將輸出的信息打印出來

5auto result = static_cast<float*>(y->data);

6for (int i = 0; i < 3; i++)

7    cout<<result[i]<<endl;

最后的輸出信息是

113.8204

2-7.31387

3-6.8253

可以看到，成功識別出了布這張圖片，到底為止在C++端的部署就完畢了。

在樹莓派上的部署

在樹莓派上的部署其實也是很簡單的，與上述步驟中不同的地方是需要設置target為樹莓派專用:

1target = tvm.target.arm_cpu('rasp3b')

點進去其實可以發現rasp3b對應着-target=armv7l-linux-gnueabihf：

 1trans_table = {

 2    "pixel2":    ["-model=snapdragon835", "-target=arm64-linux-android -mattr=+neon"],

 3    "mate10":    ["-model=kirin970", "-target=arm64-linux-android -mattr=+neon"],

 4    "mate10pro": ["-model=kirin970", "-target=arm64-linux-android -mattr=+neon"],

 5    "p20":       ["-model=kirin970", "-target=arm64-linux-android -mattr=+neon"],

 6    "p20pro":    ["-model=kirin970", "-target=arm64-linux-android -mattr=+neon"],

 7    "rasp3b":    ["-model=bcm2837", "-target=armv7l-linux-gnueabihf -mattr=+neon"],

 8    "rk3399":    ["-model=rk3399", "-target=aarch64-linux-gnu -mattr=+neon"],

 9    "pynq":      ["-model=pynq", "-target=armv7a-linux-eabi -mattr=+neon"],

10    "ultra96":   ["-model=ultra96", "-target=aarch64-linux-gnu -mattr=+neon"],

11}

還有一點改動的是，在導出.so的時候需要加入cc="/usr/bin/arm-linux-gnueabihf-g++"，此時的/usr/bin/arm-linux-gnueabihf-g++為之前下載的交叉編譯器。

1path_lib = '../tvm/deploy_lib.so'

2lib.export_library(path_lib, cc="/usr/bin/arm-linux-gnueabihf-g++")

這時就可以導出來樹莓派需要的幾個文件，之后將這幾個文件移到樹莓派中，隨后利用上面說到的C++部署代碼去部署就可以了。

關心的問題

看到這里想必大家應該還有很多疑惑，限於篇幅(寫的有點累呀)，這里講幾個比較重點的東西：

速度

這里可以毫不猶豫地說，對於這個模型來說，速度提升很明顯。在PC端部署中，使用TVM部署的手勢檢測模型的運行速度是libtorch中的5倍左右，精度還沒有測試，但是在用攝像頭進行演示過程中並沒有發現明顯的區別。當然還需要進一步的測試，就不在這里多說了。

在樹莓派中，這個模型還沒有達到實時(53ms)，但是無論對TVM，依然還有很大的優化空間，實時只是時間關系。

層的支持程度

當然因為TVM還處於開發階段，有一些層時不支持的，上文中的mobilenetv2-128_S.onnx模型一開始使用Relay IR前端讀取的時候提示，TVM中沒有flatten層的支持，mobilenetv2-128_S.onnx中有一個flatten層，所以提示報錯。

但是這個是問題嗎？只要仔細看看TVM的源碼，熟悉熟悉結構，就可以自己加層了，但其實flatten的操作函數在TVM中已經存在了，只是ONNX的前端接口沒有展示出來，onnx前端展示的是batch_flatten這個函數，其實batch_flatten就是flatten的特殊版，於是簡單修改源碼，重新編譯一下就可以成功讀取自己的模型了。

 

參考文獻

https://oldpan.me/archives/the-first-step-towards-tvm-1

https://mp.weixin.qq.com/s?__biz=Mzg3ODU2MzY5MA==&mid=2247484929&idx=1&sn=3fcce36b5a50cd8571cf932a23083667&chksm=cf109e04f86717129c3381ebeec2d0c1f7baf6ed057c66310662f5935beea88baf23e99898f4&token=1276531538&lang=zh_CN#rd

https://mp.weixin.qq.com/s?__biz=Mzg3ODU2MzY5MA==&mid=2247484930&idx=1&sn=ddc3da7b72c900ce2f8e6aad99a9e788&source=41#wechat_redirect