Permalink
Branch: master
Find file Copy path
770f6ff Mar 12, 2018
2 contributors

Users who have contributed to this file

@tqchen @phisiart
504 lines (410 sloc) 16 KB
"""
Quick Start - End-to-End Tutorial for NNVM/TVM Pipeline for OpenGL and WebGL
============================================================================
**Author**: `Zhixun Tan <https://github.com/phisiart>`_
This example shows how to build a neural network with NNVM python frontend and
generate runtime library for WebGL running in a browser with TVM. (Thanks to
Tianqi's `tutorial for cuda <http://nnvm.tvmlang.org/tutorials/get_started.html>`_ and
Ziheng's `tutorial for Raspberry Pi <http://nnvm.tvmlang.org/tutorials/deploy_model_on_rasp.html>`_)
To run this notebook, you need to install tvm and nnvm following
`these instructions <https://github.com/dmlc/nnvm/blob/master/docs/how_to/install.md>`_.
Notice that you need to build tvm with OpenGL.
"""
######################################################################
# Overview
# --------
# In this tutorial, we will download a pre-trained resnet18 model from Gluon
# Model Zoo, and run image classification in 3 different ways:
#
# - Run locally:
# We will compile the model into a TVM library with OpenGL device code and
# directly run it locally.
#
# - Run in a browser through RPC:
# We will compile the model into a JavaScript TVM library with WebGL device
# code, and upload it to an RPC server that is hosting JavaScript TVM runtime
# to run it.
#
# - Export a JavaScript library and run in a browser:
# We will compile the model into a JavaScript TVM library with WebGL device
# code, combine it with JavaScript TVM runtime, and pack everything together.
# Then we will run it directly in a browser.
#
from __future__ import print_function
import numpy as np
import tvm
import nnvm.compiler
import nnvm.testing
# This tutorial must be run with OpenGL backend enabled in TVM.
# The NNVM CI does not enable OpenGL yet. But the user can run this script.
opengl_enabled = tvm.module.enabled("opengl")
# To run the local demo, set this flag to True.
run_deploy_local = False
# To run the RPC demo, set this flag to True.
run_deploy_rpc = False
# To run the WebGL deploy demo, set this flag to True.
run_deploy_web = True
######################################################################
# Download a Pre-trained Resnet18 Model
# -------------------------------------
# Here we define 2 functions:
#
# - A function that downloads a pre-trained resnet18 model from Gluon Model Zoo.
# The model that we download is in MXNet format, we then transform it into an
# NNVM computation graph.
#
# - A function that downloads a file that contains the name of all the image
# classes in this model.
#
def load_mxnet_resnet():
"""Load a pretrained resnet model from MXNet and transform that into NNVM
format.
Returns
-------
net : nnvm.Symbol
The loaded resnet computation graph.
params : dict[str -> NDArray]
The pretrained model parameters.
data_shape: tuple
The shape of the input tensor (an image).
out_shape: tuple
The shape of the output tensor (probability of all classes).
"""
print("Loading pretrained resnet model from MXNet...")
# Download a pre-trained mxnet resnet18_v1 model.
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
# Transform the mxnet model into NNVM.
# We want a probability so add a softmax operator.
sym, params = nnvm.frontend.from_mxnet(block)
sym = nnvm.sym.softmax(sym)
print("- Model loaded!")
return sym, params, (1, 3, 224, 224), (1, 1000)
def download_synset():
"""Download a dictionary from class index to name.
This lets us know what our prediction actually is.
Returns
-------
synset : dict[int -> str]
The loaded synset.
"""
print("Downloading synset...")
from mxnet import gluon
url = "https://gist.githubusercontent.com/zhreshold/" + \
"4d0b62f3d01426887599d4f7ede23ee5/raw/" + \
"596b27d23537e5a1b5751d2b0481ef172f58b539/" + \
"imagenet1000_clsid_to_human.txt"
file_name = "synset.txt"
gluon.utils.download(url, file_name)
with open(file_name) as f:
synset = eval(f.read())
print("- Synset downloaded!")
return synset
######################################################################
# Download Input Image
# --------------------
# Here we define 2 functions that prepare an image that we want to perform
# classification on.
#
# - A function that downloads a cat image.
#
# - A function that performs preprocessing to an image so that it fits the
# format required by the resnet18 model.
#
def download_image():
"""Download a cat image and resize it to 224x224 which fits resnet.
Returns
-------
image : PIL.Image.Image
The loaded and resized image.
"""
print("Downloading cat image...")
from matplotlib import pyplot as plt
from mxnet import gluon
from PIL import Image
url = "https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true"
img_name = "cat.jpg"
gluon.utils.download(url, img_name)
image = Image.open(img_name).resize((224, 224))
print("- Cat image downloaded!")
plt.imshow(image)
plt.show()
return image
def transform_image(image):
"""Perform necessary preprocessing to input image.
Parameters
----------
image : numpy.ndarray
The raw image.
Returns
-------
image : numpy.ndarray
The preprocessed 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
######################################################################
# Compile the Model
# -----------------
# Here we define a function that invokes the NNVM compiler.
#
def compile_net(net, target_host, target, data_shape, params):
"""Compiles an NNVM computation graph.
Parameters
----------
net : nnvm.Graph
The NNVM computation graph.
target_host : str
The target to compile the host portion of the library.
target : str
The target to compile the device portion of the library.
data_shape : tuple
The shape of the input data (image).
params : dict[str -> NDArray]
Model parameters.
Returns
-------
graph : Graph
The final execution graph.
libmod : tvm.Module
The module that comes with the execution graph
params : dict[str -> NDArray]
The updated parameters of graph if params is passed.
This can be different from the params passed in.
"""
print("Compiling the neural network...")
with nnvm.compiler.build_config(opt_level=0):
deploy_graph, lib, deploy_params = nnvm.compiler.build(
net,
target_host=target_host,
target=target,
shape={"data": data_shape},
params=params)
print("- Complilation completed!")
return deploy_graph, lib, deploy_params
######################################################################
# Demo 1: Deploy Locally
# ----------------------
# In this demo, we will compile the model targetting the local machine.
#
# Then we will demonstrate how to save the compiled model as a shared library
# and load it back.
#
# Finally, we will run the model.
#
def deploy_local():
"""Runs the demo that deploys a model locally.
"""
# Load resnet model.
net, params, data_shape, out_shape = load_mxnet_resnet()
# Compile the model.
# Note that we specify the the host target as "llvm".
deploy_graph, lib, deploy_params = compile_net(
net,
target_host="llvm",
target="opengl",
data_shape=data_shape,
params=params)
# Save the compiled module.
# Note we need to save all three files returned from the NNVM compiler.
print("Saving the compiled module...")
from tvm.contrib import util
temp = util.tempdir()
path_lib = temp.relpath("deploy_lib.so")
path_graph_json = temp.relpath("deploy_graph.json")
path_params = temp.relpath("deploy_param.params")
lib.export_library(path_lib)
with open(path_graph_json, "w") as fo:
fo.write(deploy_graph.json())
with open(path_params, "wb") as fo:
fo.write(nnvm.compiler.save_param_dict(deploy_params))
print("- Saved files:", temp.listdir())
# Load the module back.
print("Loading the module back...")
loaded_lib = tvm.module.load(path_lib)
with open(path_graph_json) as fi:
loaded_graph_json = fi.read()
with open(path_params, "rb") as fi:
loaded_params = bytearray(fi.read())
print("- Module loaded!")
# Run the model! We will perform prediction on an image.
print("Running the graph...")
from tvm.contrib import graph_runtime
module = graph_runtime.create(loaded_graph_json, loaded_lib, tvm.opengl(0))
module.load_params(loaded_params)
image = transform_image(download_image())
input_data = tvm.nd.array(image.astype("float32"), ctx=tvm.opengl(0))
module.set_input("data", input_data)
module.run()
# Retrieve the output.
out = module.get_output(0, tvm.nd.empty(out_shape, ctx=tvm.opengl(0)))
top1 = np.argmax(out.asnumpy())
synset = download_synset()
print('TVM prediction top-1:', top1, synset[top1])
if run_deploy_local and opengl_enabled:
deploy_local()
######################################################################
# Demo 2: Deploy the Model to WebGL Remotely with RPC
# -------------------------------------------------------
# Following the steps above, we can also compile the model for WebGL.
# TVM provides rpc module to help with remote deploying.
#
# When we deploy a model locally to OpenGL, the model consists of two parts:
# the host LLVM part and the device GLSL part. Now that we want to deploy to
# WebGL, we need to leverage Emscripten to transform LLVM into JavaScript. In
# order to do that, we will need to specify the host target as
# 'llvm -target=asmjs-unknown-emscripten -system-lib`. Then call Emscripten to
# compile the LLVM binary output into a JavaScript file.
#
# First, we need to manually start an RPC server. Please follow the instructions
# in `tvm/web/README.md`. After following the steps, you should have a web page
# opened in a browser, and a Python script running a proxy.
#
def deploy_rpc():
"""Runs the demo that deploys a model remotely through RPC.
"""
from tvm.contrib import rpc, util, emscripten
# As usual, load the resnet18 model.
net, params, data_shape, out_shape = load_mxnet_resnet()
# Compile the model.
# Note that this time we are changing the target.
# This is because we want to translate the host library into JavaScript
# through Emscripten.
graph, lib, params = compile_net(
net,
target_host="llvm -target=asmjs-unknown-emscripten -system-lib",
target="opengl",
data_shape=data_shape,
params=params)
# Now we want to deploy our model through RPC.
# First we ned to prepare the module files locally.
print("Saving the compiled module...")
temp = util.tempdir()
path_obj = temp.relpath("deploy.bc") # host LLVM part
path_dso = temp.relpath("deploy.js") # host JavaScript part
path_gl = temp.relpath("deploy.gl") # device GLSL part
path_json = temp.relpath("deploy.tvm_meta.json")
lib.save(path_obj)
emscripten.create_js(path_dso, path_obj, side_module=True)
lib.imported_modules[0].save(path_gl)
print("- Saved files:", temp.listdir())
# Connect to the RPC server.
print("Connecting to RPC server...")
proxy_host = 'localhost'
proxy_port = 9090
remote = rpc.connect(proxy_host, proxy_port, key="js")
print("- Connected to RPC server!")
# Upload module to RPC server.
print("Uploading module to RPC server...")
remote.upload(path_dso, "deploy.dso")
remote.upload(path_gl)
remote.upload(path_json)
print("- Upload completed!")
# Load remote library.
print("Loading remote library...")
fdev = remote.load_module("deploy.gl")
fhost = remote.load_module("deploy.dso")
fhost.import_module(fdev)
rlib = fhost
print("- Remote library loaded!")
ctx = remote.opengl(0)
# Upload the parameters.
print("Uploading parameters...")
rparams = {k: tvm.nd.array(v, ctx) for k, v in params.items()}
print("- Parameters uploaded!")
# Create the remote runtime module.
print("Running remote module...")
from tvm.contrib import graph_runtime
module = graph_runtime.create(graph, rlib, ctx)
# Set parameter.
module.set_input(**rparams)
# Set input data.
input_data = np.random.uniform(size=data_shape)
module.set_input('data', tvm.nd.array(input_data.astype('float32')))
# Run.
module.run()
print("- Remote module execution completed!")
out = module.get_output(0, out=tvm.nd.empty(out_shape, ctx=ctx))
# Print first 10 elements of output.
print(out.asnumpy()[0][0:10])
if run_deploy_rpc and opengl_enabled:
deploy_rpc()
######################################################################
# Demo 3: Deploy the Model to WebGL SystemLib
# -----------------------------------------------
# This time we are not using RPC. Instead, we will compile the model and link it
# with the entire tvm runtime into a single giant JavaScript file. Then we will
# run the model using JavaScript.
#
def deploy_web():
"""Runs the demo that deploys to web.
"""
import base64
import json
import os
import shutil
import SimpleHTTPServer, SocketServer
from tvm.contrib import emscripten
curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
working_dir = os.getcwd()
output_dir = os.path.join(working_dir, "resnet")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# As usual, load the resnet18 model.
net, params, data_shape, out_shape = load_mxnet_resnet()
# As usual, compile the model.
graph, lib, params = compile_net(
net,
target_host="llvm -target=asmjs-unknown-emscripten -system-lib",
target="opengl",
data_shape=data_shape,
params=params)
# Now we save the model and link it with the TVM web runtime.
path_lib = os.path.join(output_dir, "resnet.js")
path_graph = os.path.join(output_dir, "resnet.json")
path_params = os.path.join(output_dir, "resnet.params")
path_data_shape = os.path.join(output_dir, "data_shape.json")
path_out_shape = os.path.join(output_dir, "out_shape.json")
lib.export_library(path_lib, emscripten.create_js, options=[
"-s", "USE_GLFW=3",
"-s", "USE_WEBGL2=1",
"-lglfw",
"-s", "TOTAL_MEMORY=1073741824",
])
with open(path_graph, "w") as fo:
fo.write(graph.json())
with open(path_params, "w") as fo:
fo.write(base64.b64encode(nnvm.compiler.save_param_dict(params)))
shutil.copyfile(os.path.join(curr_path, "../tvm/web/tvm_runtime.js"),
os.path.join(output_dir, "tvm_runtime.js"))
shutil.copyfile(os.path.join(curr_path, "web/resnet.html"),
os.path.join(output_dir, "resnet.html"))
# Now we want to save some extra files so that we can execute the model from
# JavaScript.
# - data shape
with open(path_data_shape, "w") as fo:
json.dump(list(data_shape), fo)
# - out shape
with open(path_out_shape, "w") as fo:
json.dump(list(out_shape), fo)
# - input image
image = download_image()
image.save(os.path.join(output_dir, "data.png"))
# - synset
synset = download_synset()
with open(os.path.join(output_dir, "synset.json"), "w") as fo:
json.dump(synset, fo)
print("Output files are in", output_dir)
# Finally, we fire up a simple web server to serve all the exported files.
print("Now running a simple server to serve the files...")
os.chdir(output_dir)
port = 8080
handler = SimpleHTTPServer.SimpleHTTPRequestHandler
httpd = SocketServer.TCPServer(("", port), handler)
print("Please open http://localhost:" + str(port) + "/resnet.html")
httpd.serve_forever()
if run_deploy_web and opengl_enabled:
deploy_web()