TVM：在樹莓派上部署預(yù)訓(xùn)練的模型

之前我們已經(jīng)介紹如何通過Python接口（AutoTVM）來編譯和優(yōu)化模型。本文將介紹如何在遠(yuǎn)程（如本例中的樹莓派）上部署預(yù)訓(xùn)練的模型。

在設(shè)備上構(gòu)建 TVM Runtime

首先我們需要再遠(yuǎn)程設(shè)備上安裝 TVM。

注意：本節(jié)和下一節(jié)中的所有指令都應(yīng)在目標(biāo)設(shè)備上執(zhí)行，例如 樹莓派。 我們假設(shè)它運行著 Linux。

由于我們在本地機器上進(jìn)行編譯，因此遠(yuǎn)程設(shè)備僅用于運行生成的代碼。 我們只需要在遠(yuǎn)程設(shè)備上構(gòu)建 TVM Runtime。這里在我們安裝 tvm 的文章中已經(jīng)詳細(xì)介紹過了，不再贅述。需要提一下的是建議在樹莓派上也安裝 llvm，并在安裝 tvm 時指定對應(yīng)的路徑。

可參考：TVM：源碼編譯安裝

git clone --recursive https://github.com/apache/tvm tvm cd tvm mkdir build cp cmake/config.cmake build cd build cmake .. make runtime -j4

構(gòu)建成功后配置 Python 包路徑：

export PYTHONPATH=$PYTHONPATH:/path/to/tvm/python

在遠(yuǎn)程設(shè)備上設(shè)置 RPC 服務(wù)器

在遠(yuǎn)程設(shè)備（如本例中的樹莓派）上運行以下命令來開啟 RPC 服務(wù)器：

python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090

如果看到下面這行說明遠(yuǎn)程設(shè)備上的 RPC 服務(wù)已經(jīng)成功開啟了：

INFO:root:RPCServer: bind to 0.0.0.0:9090

準(zhǔn)備預(yù)訓(xùn)練的模型

回到我們的主機（其上應(yīng)該安裝了完整的 TVM），不同與之前使用onnx模型，這次我們使用mxnet的模型（起始都一樣，這里只是介紹另外一種模型來源）。

不同于之前在本機上的程序，我們這里通過 MXNet 來加載模型。

加載模型、測試圖像和 sysnet（用來將 ImageNet 的標(biāo)簽轉(zhuǎn)換成類別名）：

from mxnet.gluon.model_zoo.vision import get_model from PIL import Image import numpy as np# 一行即獲取模型 block = get_model("resnet18_v1", pretrained=True)# 獲取圖像并做相應(yīng)轉(zhuǎn)換 img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true" img_name = "cat.png" img_path = download_testdata(img_url, img_name, module="data") image = Image.open(img_path).resize((224, 224))def transform_image(image):image = np.array(image) - np.array([123.0, 117.0, 104.0])image /= np.array([58.395, 57.12, 57.375])image = image.transpose((2, 0, 1))image = image[np.newaxis, :]return imagex = transform_image(image)# sysnet 用來將 ImageNet 的類別標(biāo)簽轉(zhuǎn)換為類別名稱 synset_url = "".join(["https://gist.githubusercontent.com/zhreshold/","4d0b62f3d01426887599d4f7ede23ee5/raw/","596b27d23537e5a1b5751d2b0481ef172f58b539/","imagenet1000_clsid_to_human.txt",] ) synset_name = "imagenet1000_clsid_to_human.txt" synset_path = download_testdata(synset_url, synset_name, module="data") with open(synset_path) as f:synset = eval(f.read())

現(xiàn)在我們將 Gluon 模型移植到可移植的計算圖上：

# 我們在 mxnet.gluon 中支持 MXNet 靜態(tài)圖（符號）和 HybridBlock shape_dict = {"data": x.shape} mod, params = relay.frontend.from_mxnet(block, shape_dict) # 我們要的是概率值，所以加上 softmax 算子 func = mod["main"] func = relay.Function(func.params, relay.nn.softmax(func.body), None, func.type_params, func.attrs)

以下是一些基本的數(shù)據(jù)工作負(fù)載配置：

batch_size = 1 num_classes = 1000 image_shape = (3, 224, 224) data_shape = (batch_size,) + image_shape

編譯圖

我們調(diào)用帶有圖配置和參數(shù)的 relay.build() 函數(shù)來編譯圖。 但是，我們不能在具有 ARM 指令集的設(shè)備上部署 x86 程序。 也就是說我們還需要告訴 Relay 目標(biāo)設(shè)備的編譯選項，除了 net 和 params 參數(shù)來指定深度學(xué)習(xí)工作負(fù)載。 選項也很重要，不同的選項會導(dǎo)致性能有巨大差異。

如果我們在 x86 服務(wù)器上運行示例進(jìn)行演示，直接將其設(shè)置為 llvm 即可。 如果在遠(yuǎn)程設(shè)備樹莓派上運行，我們需要指定它的指令集。 如果要在真實設(shè)備上運行本教程，請將 local_demo 設(shè)置為 False。

另外，要注意的是，官方給出的例程是針對樹莓派3B，其指令集是ARMv7，而這里如果是樹莓派4B，需要修改 target 為 aarch64。即在下面改為 target = 'llvm -mtriple=aarch64-linux-gnu'

local_demo = Trueif local_demo:target = tvm.target.Target("llvm") else:target = tvm.target.arm_cpu("rasp3b")# 上面這行是下面這行的簡單形式# target = tvm.target.Target('llvm -device=arm_cpu -model=bcm2837 -mtriple=armv7l-linux-gnueabihf -mattr=+neon')with tvm.transform.PassContext(opt_level=3):lib = relay.build(func, target, params=params)# 在 `relay.build` 之后， 我們將得到三個返回值：圖（graph），庫（library）和新的參數(shù)（new parameter）, 因為我們需要做一些優(yōu)化，這可能會改變某些參數(shù)但是保證模型的輸出結(jié)果不變# 將庫（library）保存為本地臨時目錄 tmp = utils.tempdir() lib_fname = tmp.relpath("net.tar") lib.export_library(lib_fname)

此處輸出：

/workspace/python/tvm/relay/build_module.py:333: DeprecationWarning: Please use input parameter mod (tvm.IRModule) instead of deprecated parameter mod (tvm.relay.function.Function)DeprecationWarning,

通過 RPC 將模型進(jìn)行遠(yuǎn)程部署

通過 RPC，我們可以將模型從本機遠(yuǎn)程部署到樹莓派上。

# 與遠(yuǎn)程設(shè)備建立 RPC會話 if local_demo:remote = rpc.LocalSession() else:# 這里是筆者ip，要改成你自己的host = "10.206.205.11"port = 9090remote = rpc.connect(host, port)# 將庫（library）上傳到遠(yuǎn)程設(shè)備并加載它 remote.upload(lib_fname) rlib = remote.load_module("net.tar")# 建立遠(yuǎn)程 runtime模塊 dev = remote.cpu(0) module = runtime.GraphModule(rlib["default"](dev)) # 設(shè)置輸入數(shù)據(jù) module.set_input("data", tvm.nd.array(x.astype("float32"))) # 運行 module.run() # 得到輸出 out = module.get_output(0) # 得到 top1 分類結(jié)果 top1 = np.argmax(out.numpy()) print("TVM prediction top-1: {}".format(synset[top1]))

此處輸出：

TVM prediction top-1: tiger cat

Ref：

https://tvm.apache.org/docs/how_to/deploy_models/deploy_model_on_rasp.html

總結(jié)

以上是生活随笔為你收集整理的TVM：在树莓派上部署预训练的模型的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。