本篇文章译自英文文档 Compile MXNet Models。
作者是 Joshua Z. Zhang，Kazutaka Morita。
更多 TVM 中文文档可拜访 →TVM 中文站。
本文将介绍如何用 Relay 部署 MXNet 模型。
首先装置 mxnet 模块，可通过 pip 疾速装置：
pip install mxnet --user或参考官网装置指南：https://mxnet.apache.org/versions/master/install/index.html
# 一些规范的导包import mxnet as mximport tvmimport tvm.relay as relayimport numpy as np从 Gluon Model Zoo 下载 Resnet18 模型
本节会下载预训练的 imagenet 模型，并对图像进行分类。
from tvm.contrib.download import download_testdatafrom mxnet.gluon.model_zoo.vision import get_modelfrom PIL import Imagefrom matplotlib import pyplot as pltblock = 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"synset_url = "".join( [ "https://gist.githubusercontent.com/zhreshold/", "4d0b62f3d01426887599d4f7ede23ee5/raw/", "596b27d23537e5a1b5751d2b0481ef172f58b539/", "imagenet1000_clsid_to_human.txt", ])synset_name = "imagenet1000_clsid_to_human.txt"img_path = download_testdata(img_url, "cat.png", module="data")synset_path = download_testdata(synset_url, synset_name, module="data")with open(synset_path) as f: synset = eval(f.read())image = Image.open(img_path).resize((224, 224))plt.imshow(image)plt.show()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)print("x", x.shape)
输入后果：
Downloading /workspace/.mxnet/models/resnet18_v1-a0666292.zip08d19deb-ddbf-4120-9643-fcfab19e7541 from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/resnet18_v1-a0666292.zip...x (1, 3, 224, 224)编译计算图
只需几行代码，即可将 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)接下来编译计算图：
target = "cuda"with tvm.transform.PassContext(opt_level=3): lib = relay.build(func, target, params=params)输入后果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead. "target_host parameter is going to be deprecated. "在 TVM 上执行可移植计算图
接下来用 TVM 重现雷同的前向计算：
from tvm.contrib import graph_executordev = tvm.cuda(0)dtype = "float32"m = graph_executor.GraphModule(lib["default"](dev))# 设置输出m.set_input("data", tvm.nd.array(x.astype(dtype)))# 执行m.run()# 失去输入tvm_output = m.get_output(0)top1 = np.argmax(tvm_output.numpy()[0])print("TVM prediction top-1:", top1, synset[top1])输入后果：
TVM prediction top-1: 282 tiger cat应用带有预训练权重的 MXNet 符号
MXNet 罕用 arg_params 和 aux_params 别离存储网络参数，上面将展现如何在现有 API 中应用这些权重：
def block2symbol(block): data = mx.sym.Variable("data") sym = block(data) args = {} auxs = {} for k, v in block.collect_params().items(): args[k] = mx.nd.array(v.data().asnumpy()) return sym, args, auxsmx_sym, args, auxs = block2symbol(block)# 通常将其保留/加载为检查点mx.model.save_checkpoint("resnet18_v1", 0, mx_sym, args, auxs)# 磁盘上有 "resnet18_v1-0000.params" 和 "resnet18_v1-symbol.json"对于一般性 MXNet 模型：
mx_sym, args, auxs = mx.model.load_checkpoint("resnet18_v1", 0)# 用雷同的 API 来获取 Relay 计算图mod, relay_params = relay.frontend.from_mxnet(mx_sym, shape_dict, arg_params=args, aux_params=auxs)# 反复雷同的步骤，用 TVM 运行这个模型下载 Python 源代码：from_mxnet.py
下载 Jupyter Notebook：from_mxnet.ipynb