| """ |
| Compile Tensorflow Models |
| ========================= |
| This article is an introductory tutorial to deploy tensorflow models with TVM. |
| |
| For us to begin with, tensorflow python module is required to be installed. |
| |
| Please refer to https://www.tensorflow.org/install |
| """ |
| |
| # tvm and nnvm |
| import nnvm |
| import tvm |
| |
| # os and numpy |
| import numpy as np |
| import os.path |
| |
| # Tensorflow imports |
| import tensorflow as tf |
| from tensorflow.core.framework import graph_pb2 |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import tensor_util |
| |
| # Tensorflow utility functions |
| import tvm.relay.testing.tf as tf_testing |
| |
| # Base location for model related files. |
| repo_base = 'https://github.com/dmlc/web-data/raw/master/tensorflow/models/InceptionV1/' |
| |
| # Test image |
| img_name = 'elephant-299.jpg' |
| image_url = os.path.join(repo_base, img_name) |
| |
| ###################################################################### |
| # Tutorials |
| # --------- |
| # .. note:: |
| # |
| # protobuf should be exported with :any:`add_shapes=True` option. |
| # Could use https://github.com/dmlc/web-data/tree/master/tensorflow/scripts/tf-to-nnvm.py |
| # to add shapes for existing models. |
| # |
| # Please refer docs/frontend/tensorflow.md for more details for various models |
| # from tensorflow. |
| |
| model_name = 'classify_image_graph_def-with_shapes.pb' |
| model_url = os.path.join(repo_base, model_name) |
| |
| # Image label map |
| map_proto = 'imagenet_2012_challenge_label_map_proto.pbtxt' |
| map_proto_url = os.path.join(repo_base, map_proto) |
| |
| # Human readable text for labels |
| lable_map = 'imagenet_synset_to_human_label_map.txt' |
| lable_map_url = os.path.join(repo_base, lable_map) |
| |
| # Target settings |
| # Use these commented settings to build for cuda. |
| #target = 'cuda' |
| #target_host = 'llvm' |
| #layout = "NCHW" |
| #ctx = tvm.gpu(0) |
| target = 'llvm' |
| target_host = 'llvm' |
| layout = None |
| ctx = tvm.cpu(0) |
| |
| ###################################################################### |
| # Download required files |
| # ----------------------- |
| # Download files listed above. |
| from mxnet.gluon.utils import download |
| |
| download(image_url, img_name) |
| download(model_url, model_name) |
| download(map_proto_url, map_proto) |
| download(lable_map_url, lable_map) |
| |
| ###################################################################### |
| # Import model |
| # ------------ |
| # Creates tensorflow graph definition from protobuf file. |
| |
| with tf.gfile.FastGFile(os.path.join("./", model_name), 'rb') as f: |
| graph_def = tf.GraphDef() |
| graph_def.ParseFromString(f.read()) |
| graph = tf.import_graph_def(graph_def, name='') |
| # Call the utility to import the graph definition into default graph. |
| graph_def = tf_testing.ProcessGraphDefParam(graph_def) |
| # Add shapes to the graph. |
| with tf.Session() as sess: |
| graph_def = tf_testing.AddShapesToGraphDef(sess, 'softmax') |
| |
| ###################################################################### |
| # Decode image |
| # ------------ |
| # .. note:: |
| # |
| # tensorflow frontend import doesn't support preprocessing ops like JpegDecode. |
| # JpegDecode is bypassed (just return source node). |
| # Hence we supply decoded frame to TVM instead. |
| # |
| |
| from PIL import Image |
| image = Image.open(img_name).resize((299, 299)) |
| |
| x = np.array(image) |
| |
| ###################################################################### |
| # Import the graph to NNVM |
| # ------------------------ |
| # Import tensorflow graph definition to nnvm. |
| # |
| # Results: |
| # sym: nnvm graph for given tensorflow protobuf. |
| # params: params converted from tensorflow params (tensor protobuf). |
| sym, params = nnvm.frontend.from_tensorflow(graph_def, layout=layout) |
| |
| print ("Tensorflow protobuf imported as nnvm graph") |
| ###################################################################### |
| # NNVM Compilation |
| # ---------------- |
| # Compile the graph to llvm target with given input specification. |
| # |
| # Results: |
| # graph: Final graph after compilation. |
| # params: final params after compilation. |
| # lib: target library which can be deployed on target with tvm runtime. |
| |
| import nnvm.compiler |
| shape_dict = {'DecodeJpeg/contents': x.shape} |
| dtype_dict = {'DecodeJpeg/contents': 'uint8'} |
| graph, lib, params = nnvm.compiler.build(sym, shape=shape_dict, target=target, target_host=target_host, dtype=dtype_dict, params=params) |
| |
| ###################################################################### |
| # Execute the portable graph on TVM |
| # --------------------------------- |
| # Now we can try deploying the NNVM compiled model on target. |
| |
| from tvm.contrib import graph_runtime |
| dtype = 'uint8' |
| m = graph_runtime.create(graph, lib, ctx) |
| # set inputs |
| m.set_input('DecodeJpeg/contents', tvm.nd.array(x.astype(dtype))) |
| m.set_input(**params) |
| # execute |
| m.run() |
| # get outputs |
| tvm_output = m.get_output(0, tvm.nd.empty(((1, 1008)), 'float32')) |
| |
| ###################################################################### |
| # Process the output |
| # ------------------ |
| # Process the model output to human readable text for InceptionV1. |
| predictions = tvm_output.asnumpy() |
| predictions = np.squeeze(predictions) |
| |
| # Creates node ID --> English string lookup. |
| node_lookup = tf_testing.NodeLookup(label_lookup_path=os.path.join("./", map_proto), |
| uid_lookup_path=os.path.join("./", lable_map)) |
| |
| # Print top 5 predictions from TVM output. |
| top_k = predictions.argsort()[-5:][::-1] |
| for node_id in top_k: |
| human_string = node_lookup.id_to_string(node_id) |
| score = predictions[node_id] |
| print('%s (score = %.5f)' % (human_string, score)) |
| |
| ###################################################################### |
| # Inference on tensorflow |
| # ----------------------- |
| # Run the corresponding model on tensorflow |
| |
| def create_graph(): |
| """Creates a graph from saved GraphDef file and returns a saver.""" |
| # Creates graph from saved graph_def.pb. |
| with tf.gfile.FastGFile(model_name, 'rb') as f: |
| graph_def = tf.GraphDef() |
| graph_def.ParseFromString(f.read()) |
| graph = tf.import_graph_def(graph_def, name='') |
| # Call the utility to import the graph definition into default graph. |
| graph_def = tf_testing.ProcessGraphDefParam(graph_def) |
| |
| def run_inference_on_image(image): |
| """Runs inference on an image. |
| |
| Parameters |
| ---------- |
| image: String |
| Image file name. |
| |
| Returns |
| ------- |
| Nothing |
| """ |
| if not tf.gfile.Exists(image): |
| tf.logging.fatal('File does not exist %s', image) |
| image_data = tf.gfile.FastGFile(image, 'rb').read() |
| |
| # Creates graph from saved GraphDef. |
| create_graph() |
| |
| with tf.Session() as sess: |
| softmax_tensor = sess.graph.get_tensor_by_name('softmax:0') |
| predictions = sess.run(softmax_tensor, |
| {'DecodeJpeg/contents:0': image_data}) |
| |
| predictions = np.squeeze(predictions) |
| |
| # Creates node ID --> English string lookup. |
| node_lookup = tf_testing.NodeLookup(label_lookup_path=os.path.join("./", map_proto), |
| uid_lookup_path=os.path.join("./", lable_map)) |
| |
| # Print top 5 predictions from tensorflow. |
| top_k = predictions.argsort()[-5:][::-1] |
| print ("===== TENSORFLOW RESULTS =======") |
| for node_id in top_k: |
| human_string = node_lookup.id_to_string(node_id) |
| score = predictions[node_id] |
| print('%s (score = %.5f)' % (human_string, score)) |
| |
| run_inference_on_image (img_name) |