Compile MXNet Models

Author: Joshua Z. Zhang

This article is an introductory tutorial to deploy mxnet models with NNVM.

For us to begin with, mxnet module is required to be installed.

A quick solution is ` pip install mxnet --user ` or please refer to offical installation guide.

# some standard imports
import mxnet as mx
import nnvm
import tvm
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 get_model
from mxnet.gluon.utils import download
from PIL import Image
from matplotlib import pyplot as plt
block = get_model('resnet18_v1', pretrained=True)
img_name = 'cat.jpg'
synset_url = ''.join(['',
synset_name = 'synset.txt'
download('', img_name)
download(synset_url, synset_name)
with open(synset_name) as f:
    synset = eval(
image =, 224))

def transform_image(image):
    image = np.array(image) - np.array([123., 117., 104.])
    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)


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

sym, params = nnvm.frontend.from_mxnet(block)
# we want a probability so add a softmax operator
sym = nnvm.sym.softmax(sym)

now compile the graph

import nnvm.compiler
target = 'cuda'
shape_dict = {'data': x.shape}
graph, lib, params =, target, shape_dict, 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_runtime
ctx = tvm.gpu(0)
dtype = 'float32'
m = graph_runtime.create(graph, lib, ctx)
# set inputs
m.set_input('data', tvm.nd.array(x.astype(dtype)))
# execute
# get outputs
tvm_output = m.get_output(0, tvm.nd.empty((1000,), dtype))
top1 = np.argmax(tvm_output.asnumpy())
print('TVM prediction top-1:', top1, synset[top1])


TVM prediction top-1: 281 tabby, tabby cat

Use MXNet symbol with pretrained weights

MXNet often use arg_prams 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 NNVM compatible symbol
nnvm_sym, nnvm_params = nnvm.frontend.from_mxnet(mx_sym, args, auxs)
# repeat the same steps to run this model using TVM

Total running time of the script: ( 0 minutes 5.144 seconds)

Generated by Sphinx-Gallery