forked from apache/tvm
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[TUTORIAL] Deploy the Pretrained Model on Raspberry Pi (apache#61)
* [TUTORIAL] Deploy the Pretrained Model on Raspberry Pi * [TUTORIAL] Improve * [TUTORIAL] Improve * [TUTORIAL] Improve * [TUTORIAL] Improve
- Loading branch information
1 parent
36eb045
commit a81b025
Showing
1 changed file
with
234 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,234 @@ | ||
""" | ||
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 | ||
it on raspberry pi. | ||
To begin with, we import nnvm(for compilation) and TVM(for deployment). | ||
""" | ||
import tvm | ||
import nnvm.compiler | ||
import nnvm.testing | ||
from tvm.contrib import util, rpc | ||
from tvm.contrib import graph_runtime as runtime | ||
|
||
|
||
###################################################################### | ||
# Build TVM Runtime on Device | ||
# --------------------------- | ||
# | ||
# There're some prerequisites: we need build tvm runtime and set up | ||
# a RPC server on remote device. | ||
# | ||
# To get started, clone tvm repo from github. It is important to clone | ||
# the submodules along, with --recursive option (Assuming you are in | ||
# your home directory): | ||
# | ||
# .. code-block:: bash | ||
# | ||
# git clone --recursive https://github.com/dmlc/tvm | ||
# | ||
# .. note:: | ||
# | ||
# Usually device has limited resources and we only need to build | ||
# runtime. The idea is we will use TVM compiler on the local server | ||
# to compile and upload the compiled program to the device and run | ||
# the device function remotely. | ||
# | ||
# .. code-block:: bash | ||
# | ||
# make runtime | ||
# | ||
# After success of buildind runtime, we need set environment varibles | ||
# in :code:`~/.bashrc` file of yourself account or :code:`/etc/profile` | ||
# of system enviroment variables. Assuming your TVM directory is in | ||
# :code:`~/tvm` and set environment variables below your account. | ||
# | ||
# .. code-block:: bash | ||
# | ||
# vi ~/.bashrc | ||
# | ||
# We need edit :code:`~/.bashrc` using :code:`vi ~/.bashrc` and add | ||
# lines below (Assuming your TVM directory is in :code:`~/tvm`): | ||
# | ||
# .. code-block:: bash | ||
# | ||
# export TVM_HOME=~/tvm | ||
# export PATH=$PATH:$TVM_HOME/lib | ||
# export PYTHONPATH=$PYTHONPATH:$TVM_HOME/python | ||
# | ||
# To enable updated :code:`~/.bashrc`, execute :code:`source ~/.bashrc`. | ||
|
||
###################################################################### | ||
# Set Up RPC Server on Device | ||
# --------------------------- | ||
# To set up a TVM RPC server on the Raspberry Pi (our remote device), | ||
# we have prepared a one-line script so you only need to run this | ||
# command after following the installation guide to install TVM on | ||
# your device: | ||
# | ||
# .. code-block:: bash | ||
# | ||
# python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090 | ||
# | ||
# After executing command above, if you see these lines below, it's | ||
# successful to start RPC server on your device. | ||
# | ||
# .. code-block:: bash | ||
# | ||
# Loading runtime library /home/YOURNAME/code/tvm/lib/libtvm_runtime.so... exec only | ||
# INFO:root:RPCServer: bind to 0.0.0.0:9090 | ||
# | ||
###################################################################### | ||
# For demonstration, we simply start an RPC server on the same machine, | ||
# if :code:`use_rasp` is False. If you have set up the remote | ||
# environment, please change the three lines below: change the | ||
# :code:`use_rasp` to True, also change the host and port with your | ||
# device's host address and port number. | ||
|
||
use_rasp = False | ||
host = 'rasp0' | ||
port = 9090 | ||
|
||
if not use_rasp: | ||
# run server locally | ||
host = 'localhost' | ||
port = 9090 | ||
server = rpc.Server(host=host, port=port) | ||
|
||
###################################################################### | ||
# Prepare the Pretrained Model | ||
# ---------------------------- | ||
# Back to the host machine, firstly, we need to download a MXNet Gluon | ||
# ResNet model from model zoo, which is pretrained on ImageNet. You | ||
# can found more details about this part at `Compile MXNet Models` | ||
|
||
from mxnet.gluon.model_zoo.vision import get_model | ||
from mxnet.gluon.utils import download | ||
from PIL import Image | ||
import numpy as np | ||
|
||
# only one line to get the model | ||
block = get_model('resnet18_v1', pretrained=True) | ||
|
||
###################################################################### | ||
# In order to test our model, here we download a image of cat and | ||
# transform its format. | ||
img_name = 'cat.jpg' | ||
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 | ||
out_shape = (batch_size, num_classes) | ||
|
||
###################################################################### | ||
# 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. | ||
|
||
###################################################################### | ||
# If we run the example locally for demonstration, we can simply set | ||
# it as :code:`llvm`. If to run it on the Raspberry Pi, you need to | ||
# specify its instruction set. Here is the option I use for my Raspberry | ||
# Pi, which has been proved as a good compilation configuration. | ||
|
||
if use_rasp: | ||
target = "llvm -target=armv7l-none-linux-anueabihf -mcpu=cortex-a53 -mattr=+neon" | ||
else: | ||
target = "llvm" | ||
|
||
|
||
# use `with tvm.target.rasp` for some target-specified optimization | ||
with tvm.target.rasp(): | ||
graph, lib, params = nnvm.compiler.build( | ||
net, target, shape={"data": data_shape}, params=params) | ||
|
||
# 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() | ||
lib_fname = tmp.relpath('net.o') | ||
lib.save(lib_fname) | ||
|
||
|
||
###################################################################### | ||
# Deploy the Model Remotely by RPC | ||
# -------------------------------- | ||
# With RPC, you can deploy the model remotely from your host machine | ||
# to the remote device. | ||
|
||
# connect the server | ||
remote = rpc.connect(host, port) | ||
|
||
# upload the library to remote device and load it | ||
remote.upload(lib_fname) | ||
rlib = remote.load_module('net.o') | ||
|
||
ctx = remote.cpu(0) | ||
# upload the parameter | ||
rparams = {k: tvm.nd.array(v, ctx) for k, v in params.items()} | ||
|
||
# create the remote runtime module | ||
module = runtime.create(graph, rlib, ctx) | ||
# set parameter | ||
module.set_input(**rparams) | ||
# set input data | ||
module.set_input('data', tvm.nd.array(x.astype('float32'))) | ||
# run | ||
module.run() | ||
# get output | ||
out = module.get_output(0, tvm.nd.empty(out_shape, ctx=ctx)) | ||
# get top1 result | ||
top1 = np.argmax(out.asnumpy()) | ||
print('TVM prediction top-1: {}'.format(synset[top1])) | ||
|
||
if not use_rasp: | ||
# terminate the local server | ||
server.terminate() |