compile-rasp3b.py

import nnvm
import tvm
import numpy as np

from mxnet.gluon.model_zoo.vision import get_model

batch_size = 1
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape
model_name = 'resnet18_v1' 

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

opt_level = 3
#target = 'llvm'
#target = 'llvm -device=arm_cpu -model=bcm2837 -target=armv7l-linux-gnueabihf -mattr=+neon'
target = tvm.target.arm_cpu('rasp3b')
print('model:', model_name, ', target:', target, ', opt_level:', opt_level, ', data_shape:', data_shape)
print("Compiling...")

with nnvm.compiler.build_config(opt_level=opt_level):
    graph, lib, params = nnvm.compiler.build(sym, target, shape={"data": data_shape}, params=params)
print("Compilation done")

print("Saving files")
# save the graph, lib and params into separate files
path_lib = "model.tar"
lib.export_library(path_lib)
with open("model.json", "w") as fo:
    fo.write(graph.json())
with open("model.params", "wb") as fo:
    fo.write(nnvm.compiler.save_param_dict(params))
print("Files saved")