deploy_model_on_rasp.py 6.97 KB
Newer Older
1
"""
2 3
.. _tutorial-deploy-model-on-rasp:

4 5 6 7 8
Deploy the Pretrained Model on Raspberry Pi
===========================================
**Author**: `Ziheng Jiang <https://ziheng.org/>`_

This is an example of using NNVM to compile a ResNet model and deploy
9
it on Raspberry Pi.
10
"""
11

12 13 14
import tvm
import nnvm.compiler
import nnvm.testing
15 16
from tvm import rpc
from tvm.contrib import util, graph_runtime as runtime
17 18

######################################################################
19 20
# .. _build-tvm-runtime-on-device:
#
21 22 23
# Build TVM Runtime on Device
# ---------------------------
#
24
# The first step is to build tvm runtime on the remote device.
25 26 27
#
# .. note::
#
28 29 30 31 32 33 34
#   All instructions in both this section and next section should be
#   executed on the target device, e.g. Raspberry Pi. And we assume it
#   has Linux running.
# 
# Since we do compilation on local machine, the remote device is only used
# for running the generated code. We only need to build tvm runtime on
# the remote device.
35
#
36
# .. code-block:: bash
37
#
38 39 40
#   git clone --recursive https://github.com/dmlc/tvm
#   cd tvm
#   make runtime -j4
41
#
42 43 44 45
# After building runtime successfully, we need to set environment varibles
# in :code:`~/.bashrc` file. We can edit :code:`~/.bashrc`
# using :code:`vi ~/.bashrc` and add the line below (Assuming your TVM 
# directory is in :code:`~/tvm`):
46
#
47
# .. code-block:: bash
48
#
49
#   export PYTHONPATH=$PYTHONPATH:~/tvm/python
50
#
51
# To update the environment variables, execute :code:`source ~/.bashrc`.
52 53 54 55

######################################################################
# Set Up RPC Server on Device
# ---------------------------
56 57
# To start an RPC server, run the following command on your remote device
# (Which is Raspberry Pi in our example).
58 59 60 61 62
#
#   .. code-block:: bash
#
#     python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
#
63 64
# If you see the line below, it means the RPC server started
# successfully on your device.
65 66 67 68
#
#    .. code-block:: bash
#
#      INFO:root:RPCServer: bind to 0.0.0.0:9090
69
#
70 71

######################################################################
72 73 74 75 76 77 78
# Prepare the Pre-trained Model
# -----------------------------
# Back to the host machine, which should have a full TVM installed (with LLVM).
# 
# We will use pre-trained model from
# `MXNet Gluon model zoo <https://mxnet.incubator.apache.org/api/python/gluon/model_zoo.html>`_.
# You can found more details about this part at tutorial :ref:`tutorial-from-mxnet`.
79 80 81 82 83 84

from mxnet.gluon.model_zoo.vision import get_model
from mxnet.gluon.utils import download
from PIL import Image
import numpy as np

85
# one line to get the model
86 87 88
block = get_model('resnet18_v1', pretrained=True)

######################################################################
89
# In order to test our model, here we download an image of cat and
90
# transform its format.
91
img_name = 'cat.png'
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
download('https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true', img_name)
image = Image.open(img_name).resize((224, 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)

######################################################################
# synset is used to transform the label from number of ImageNet class to
# the word human can understand.
synset_url = ''.join(['https://gist.githubusercontent.com/zhreshold/',
                      '4d0b62f3d01426887599d4f7ede23ee5/raw/',
                      '596b27d23537e5a1b5751d2b0481ef172f58b539/',
                      'imagenet1000_clsid_to_human.txt'])
synset_name = 'synset.txt'
download(synset_url, synset_name)
with open(synset_name) as f:
    synset = eval(f.read())

######################################################################
# 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
net, params = nnvm.frontend.from_mxnet(block)
# we want a probability so add a softmax operator
net = nnvm.sym.softmax(net)

######################################################################
# Here are some basic data workload configurations.
batch_size = 1
num_classes = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape

######################################################################
# Compile The Graph
# -----------------
# To compile the graph, we call the :any:`nnvm.compiler.build` function
# with the graph configuration and parameters. However, You cannot to
# deploy a x86 program on a device with ARM instruction set. It means
# NNVM also needs to know the compilation option of target device,
# apart from arguments :code:`net` and :code:`params` to specify the
# deep learning workload. Actually, the option matters, different option
# will lead to very different performance.

######################################################################
144 145 146 147
# If we run the example on our x86 server for demonstration, we can simply
# set it as :code:`llvm`. If running it on the Raspberry Pi, we need to
# specify its instruction set. Set :code:`local_demo` to False if you want
# to run this tutorial with a real device.
148

149 150 151
local_demo = True

if local_demo:
152
    target = tvm.target.create('llvm')
153 154 155
else:
    target = tvm.target.arm_cpu('rasp3b')
    # The above line is a simple form of
156
    # target = tvm.target.create('llvm -device=arm_cpu -model=bcm2837 -target=armv7l-linux-gnueabihf -mattr=+neon')
157

158
with nnvm.compiler.build_config(opt_level=3):
159 160
    graph, lib, params = nnvm.compiler.build(
        net, target, shape={"data": data_shape}, params=params)
161 162 163 164 165 166 167

# After `nnvm.compiler.build`, you will get three return values: graph,
# library and the new parameter, since we do some optimization that will
# change the parameters but keep the result of model as the same.

# Save the library at local temporary directory.
tmp = util.tempdir()
168 169
lib_fname = tmp.relpath('net.tar')
lib.export_library(lib_fname)
170 171 172 173 174 175 176

######################################################################
# Deploy the Model Remotely by RPC
# --------------------------------
# With RPC, you can deploy the model remotely from your host machine
# to the remote device.

177 178 179 180 181 182 183 184
# obtain an RPC session from remote device.
if local_demo:
    remote = rpc.LocalSession()
else:
    # The following is my environment, change this to the IP address of your target device
    host = '10.77.1.162'
    port = 9090
    remote = rpc.connect(host, port)
185 186 187

# upload the library to remote device and load it
remote.upload(lib_fname)
188
rlib = remote.load_module('net.tar')
189 190

# create the remote runtime module
191
ctx = remote.cpu(0)
192
module = runtime.create(graph, rlib, ctx)
193 194
# set parameter (upload params to the remote device. This may take a while)
module.set_input(**params)
195 196 197 198 199
# set input data
module.set_input('data', tvm.nd.array(x.astype('float32')))
# run
module.run()
# get output
200
out = module.get_output(0)
201 202 203
# get top1 result
top1 = np.argmax(out.asnumpy())
print('TVM prediction top-1: {}'.format(synset[top1]))