Compile MXNet Models

Author: Joshua Z. Zhang, Kazutaka Morita

This article is an introductory tutorial to deploy mxnet models with Relay. To begin, we must install mxnet:

pip install mxnet

or please refer to official installation guide.

# some standard imports
import mxnet as mx
import tvm
import tvm.relay as relay
import numpy as np

Download Resnet18 model from Gluon Model Zoo

In this section, we download a pretrained imagenet model and classify an image.

from import download_testdata
from import get_model
from PIL import Image
from matplotlib import pyplot as plt

block = get_model("resnet18_v1", pretrained=True)
img_url = ""
img_name = "cat.png"
synset_url = "".join(
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(
image =, 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)
print("x", x.shape)
from mxnet
Downloading /workspace/.mxnet/models/resnet18_v1-a0666292.zipb7ec27d1-a156-48a2-b266-665071bc2c8e from
x (1, 3, 224, 224)

Compile the Graph

Now we would like to port the Gluon model to a portable computational graph. It’s as easy as several lines. We support MXNet static graph(symbol) and HybridBlock in mxnet.gluon

shape_dict = {"data": x.shape}
mod, params = relay.frontend.from_mxnet(block, shape_dict)
## we want a probability so add a softmax operator
func = mod["main"]
func = relay.Function(func.params, relay.nn.softmax(func.body), None, func.type_params, func.attrs)

now compile the graph

target = "cuda"
with tvm.transform.PassContext(opt_level=3):
    lib =, target, params=params)

Execute the portable graph on TVM

Now, we would like to reproduce the same forward computation using TVM.

from tvm.contrib import graph_executor

dev = tvm.cuda(0)
dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# set inputs
m.set_input("data", tvm.nd.array(x.astype(dtype)))
# execute
# get outputs
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

Use MXNet symbol with pretrained weights

MXNet often use arg_params and aux_params to store network parameters separately, here we show how to use these weights with existing 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(
    return sym, args, auxs

mx_sym, args, auxs = block2symbol(block)
# usually we would save/load it as checkpoint
mx.model.save_checkpoint("resnet18_v1", 0, mx_sym, args, auxs)
# there are 'resnet18_v1-0000.params' and 'resnet18_v1-symbol.json' on disk

for a normal mxnet model, we start from here

mx_sym, args, auxs = mx.model.load_checkpoint("resnet18_v1", 0)
# now we use the same API to get Relay computation graph
mod, relay_params = relay.frontend.from_mxnet(mx_sym, shape_dict, arg_params=args, aux_params=auxs)
# repeat the same steps to run this model using TVM

Gallery generated by Sphinx-Gallery