tutorials/nnvm/from_mxnet_to_webgl.py

"""
Deploy Deep Learning Models to 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.
To run this notebook, you need to install tvm and nnvm.
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 = False

######################################################################
# 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.png"

    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 import rpc
    from tvm.contrib import 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(os.getcwd())))
    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()