TVM：在树莓派上部署预训练的模型

TVM：在树莓派上部署预训练的模型

之前我们已经介绍如何通过Python接口（AutoTVM）来编译和优化模型。本文将介绍如何在远程（如本例中的树莓派）上部署预训练的模型。

在设备上构建 TVM Runtime

首先我们需要再远程设备上安装 TVM。

注意：本节和下一节中的所有指令都应在目标设备上执行，例如 树莓派。 我们假设它运行着 Linux。

由于我们在本地机器上进行编译，因此远程设备仅用于运行生成的代码。 我们只需要在远程设备上构建 TVM Runtime。这里在我们安装 tvm 的文章中已经详细介绍过了，不再赘述。需要提一下的是建议在树莓派上也安装 llvm，并在安装 tvm 时指定对应的路径。

可参考：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

构建成功后配置 Python 包路径：

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

在远程设备上设置 RPC 服务器

在远程设备（如本例中的树莓派）上运行以下命令来开启 RPC 服务器：

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

如果看到下面这行说明远程设备上的 RPC 服务已经成功开启了：

INFO:root:RPCServer: bind to 0.0.0.0:9090

准备预训练的模型

回到我们的主机（其上应该安装了完整的 TVM），不同与之前使用onnx模型，这次我们使用mxnet的模型（起始都一样，这里只是介绍另外一种模型来源）。

不同于之前在本机上的程序，我们这里通过 MXNet 来加载模型。

加载模型、测试图像和 sysnet（用来将 ImageNet 的标签转换成类别名）：

from mxnet.gluon.model_zoo.vision import get_model
from PIL import Image
import numpy as np

# 一行即获取模型
block = get_model("resnet18_v1", pretrained=True)

# 获取图像并做相应转换
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 image

x = transform_image(image)

# sysnet 用来将 ImageNet 的类别标签转换为类别名称
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())

现在我们将 Gluon 模型移植到可移植的计算图上：

# 我们在 mxnet.gluon 中支持 MXNet 静态图（符号）和 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)

以下是一些基本的数据工作负载配置：

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

编译图

我们调用带有图配置和参数的 relay.build() 函数来编译图。 但是，我们不能在具有 ARM 指令集的设备上部署 x86 程序。 也就是说我们还需要告诉 Relay 目标设备的编译选项，除了 net 和 params 参数来指定深度学习工作负载。 选项也很重要，不同的选项会导致性能有巨大差异。

如果我们在 x86 服务器上运行示例进行演示，直接将其设置为 llvm 即可。 如果在远程设备树莓派上运行，我们需要指定它的指令集。 如果要在真实设备上运行本教程，请将 local_demo 设置为 False。

另外，要注意的是，官方给出的例程是针对树莓派3B，其指令集是ARMv7，而这里如果是树莓派4B，需要修改 target 为 aarch64。即在下面改为 target = 'llvm -mtriple=aarch64-linux-gnu'

local_demo = True

if 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）和新的参数（new parameter）, 因为我们需要做一些优化，这可能会改变某些参数但是保证模型的输出结果不变

# 将库（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 将模型进行远程部署

通过 RPC，我们可以将模型从本机远程部署到树莓派上。

# 与远程设备建立 RPC会话
if local_demo:
    remote = rpc.LocalSession()
else:
    # 这里是笔者ip，要改成你自己的
    host = "10.206.205.11"
    port = 9090
    remote = rpc.connect(host, port)

# 将库（library）上传到远程设备并加载它
remote.upload(lib_fname)
rlib = remote.load_module("net.tar")

# 建立远程 runtime模块
dev = remote.cpu(0)
module = runtime.GraphModule(rlib["default"](dev))
# 设置输入数据
module.set_input("data", tvm.nd.array(x.astype("float32")))
# 运行
module.run()
# 得到输出
out = module.get_output(0)
# 得到 top1 分类结果
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