この記事は、英語のドキュメントCompile MXNet Modelsを翻訳したものです。

さらに多くの TVM 中国語ドキュメントにアクセスできます → TVM 中国語ステーション。

この記事では、Relay を使用して MXNet モデルをデプロイする方法を紹介します。

まず mxnet モジュールをインストールします。これは pip 経由ですぐにインストールできます。

pip install mxnet --user

または、公式インストールガイドを参照してください: https://mxnet.apache.org/versions/master/install/index.html

# 一些标准的导包
import mxnet as mx
import tvm
import tvm.relay as relay
import numpy as np

Gluon Model Zoo から Resnet18 モデルをダウンロードする

このセクションでは、事前トレーニングされた imagenet モデルをダウンロードし、画像を分類します。

from tvm.contrib.download import download_testdata
from mxnet.gluon.model_zoo.vision import get_model
from PIL import Image
from matplotlib import pyplot as plt

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"
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 image

x = 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 静的グラフ (シンボル) と HybridBlocks をサポートします。

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_executor

dev = 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, auxs

mx_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 ノートブックをダウンロード: from_mxnet.ipynb

MXNet モデルをコンパイルする

Gluon Model Zoo から Resnet18 モデルをダウンロードする

計算グラフを作成する

TVM でポータブルな計算グラフを実行

事前トレーニングされた重みを使用した MXNet 表記の使用

おすすめ