|
| 1 | +# some standard imports |
| 2 | +import nnvm |
| 3 | +import tvm |
| 4 | +from nnvm.compiler import graph_attr |
| 5 | +import vta |
| 6 | +import os |
| 7 | +import numpy as np |
| 8 | +from PIL import Image |
| 9 | +import pickle |
| 10 | +import json |
| 11 | +import logging |
| 12 | +import wget |
| 13 | +from tvm.contrib import graph_runtime, rpc, util |
| 14 | + |
| 15 | +factor = 16 |
| 16 | +host = "pynq" |
| 17 | +port = 9091 |
| 18 | +verbose = False |
| 19 | +# only run fpga component, mark non-conv ops as nop |
| 20 | +debug_fpga_only = False |
| 21 | + |
| 22 | +# Obtain model and hardware files (they're too large to check-in) |
| 23 | +url = "https://homes.cs.washington.edu/~moreau/media/vta/" |
| 24 | +TEST_FILE = 'cat.jpg' |
| 25 | +CATEG_FILE = 'synset.txt' |
| 26 | +RESNET_GRAPH_FILE = 'quantize_graph.json' |
| 27 | +RESNET_PARAMS_FILE = 'quantize_params.pkl' |
| 28 | +BITSTREAM_FILE = 'vta.bit' |
| 29 | +for file in [TEST_FILE, CATEG_FILE, RESNET_GRAPH_FILE, RESNET_PARAMS_FILE, BITSTREAM_FILE]: |
| 30 | + if not os.path.isfile(file): |
| 31 | + print "Downloading {}".format(file) |
| 32 | + wget.download(url+file) |
| 33 | + |
| 34 | +# Program the FPGA remotely |
| 35 | +assert tvm.module.enabled("rpc") |
| 36 | +remote = rpc.connect(host, port) |
| 37 | +remote.upload(BITSTREAM_FILE, BITSTREAM_FILE) |
| 38 | +fprogram = remote.get_function("tvm.contrib.vta.init") |
| 39 | +fprogram(BITSTREAM_FILE) |
| 40 | + |
| 41 | +if verbose: |
| 42 | + logging.basicConfig(level=logging.INFO) |
| 43 | + |
| 44 | +# Change to -device=tcpu to run cpu only inference. |
| 45 | +target = "llvm -device=vta" |
| 46 | + |
| 47 | +synset = eval(open(os.path.join(CATEG_FILE)).read()) |
| 48 | +image = Image.open(os.path.join(TEST_FILE)).resize((224, 224)) |
| 49 | + |
| 50 | +def transform_image(image): |
| 51 | + image = np.array(image) - np.array([123., 117., 104.]) |
| 52 | + image /= np.array([58.395, 57.12, 57.375]) |
| 53 | + image = image.transpose((2, 0, 1)) |
| 54 | + image = image[np.newaxis, :] |
| 55 | + return image |
| 56 | + |
| 57 | +def mark_nop(graph, conv_layer=-1, skip_conv_layer=()): |
| 58 | + """Helper function to mark certain op as nop |
| 59 | +
|
| 60 | + Useful to debug performance issues. |
| 61 | + """ |
| 62 | + jgraph = json.loads(graph.json()) |
| 63 | + counter = 0 |
| 64 | + for nid, node in enumerate(jgraph["nodes"]): |
| 65 | + op_name = node["op"] |
| 66 | + if op_name != "tvm_op": |
| 67 | + continue |
| 68 | + attrs = node["attrs"] |
| 69 | + node_name = node["name"] |
| 70 | + func_name = attrs["func_name"] |
| 71 | + if func_name.find("quantized_conv2d") != -1: |
| 72 | + if conv_layer >= 0: |
| 73 | + if counter != conv_layer: |
| 74 | + attrs["func_name"] = "__nop" |
| 75 | + if counter in skip_conv_layer: |
| 76 | + attrs["func_name"] = "__nop" |
| 77 | + counter += 1 |
| 78 | + else: |
| 79 | + if conv_layer >= 0: |
| 80 | + attrs["func_name"] = "__nop" |
| 81 | + attrs["func_name"] = "__nop" |
| 82 | + if attrs["func_name"] != "__nop": |
| 83 | + print("Run function %s"% func_name) |
| 84 | + graph = nnvm.graph.load_json(json.dumps(jgraph)) |
| 85 | + return graph |
| 86 | + |
| 87 | +x = transform_image(image) |
| 88 | +print('x', x.shape) |
| 89 | + |
| 90 | +###################################################################### |
| 91 | +# now compile the graph |
| 92 | +import nnvm.compiler |
| 93 | +np.random.seed(0) |
| 94 | +sym = nnvm.graph.load_json( |
| 95 | + open(os.path.join(RESNET_GRAPH_FILE)).read()) |
| 96 | +params = pickle.load( |
| 97 | + open(os.path.join(RESNET_PARAMS_FILE))) |
| 98 | + |
| 99 | +shape_dict = {"data": x.shape} |
| 100 | +dtype_dict = {"data": 'float32'} |
| 101 | +shape_dict.update({k: v.shape for k, v in params.items()}) |
| 102 | +dtype_dict.update({k: str(v.dtype) for k, v in params.items()}) |
| 103 | + |
| 104 | +graph = nnvm.graph.create(sym) |
| 105 | +graph_attr.set_shape_inputs(sym, shape_dict) |
| 106 | +graph_attr.set_dtype_inputs(sym, dtype_dict) |
| 107 | +graph = graph.apply("InferShape").apply("InferType") |
| 108 | + |
| 109 | +dtype = "float32" |
| 110 | +sym = vta.graph.remove_stochastic(sym) |
| 111 | +sym = vta.graph.clean_cast(sym) |
| 112 | +sym = vta.graph.clean_conv_fuse(sym) |
| 113 | +if "vta" in target: |
| 114 | + sym = vta.graph.pack(sym, shape_dict, factor) |
| 115 | + |
| 116 | +graph_attr.set_shape_inputs(sym, shape_dict) |
| 117 | +sym = sym.apply("InferShape") |
| 118 | +graph_attr.set_dtype_inputs(sym, dtype_dict) |
| 119 | +sym = sym.apply("InferType") |
| 120 | + |
| 121 | +with nnvm.compiler.build_config(opt_level=3): |
| 122 | + bdict = {} |
| 123 | + if "vta" not in target: |
| 124 | + bdict = {"add_lower_pass": []} |
| 125 | + else: |
| 126 | + bdict = {"add_lower_pass": vta.debug_mode(0)} |
| 127 | + with tvm.build_config(**bdict): |
| 128 | + graph, lib, params = nnvm.compiler.build( |
| 129 | + sym, target, shape_dict, dtype_dict, |
| 130 | + params=params) |
| 131 | + |
| 132 | +remote = rpc.connect(host, port) |
| 133 | +temp = util.tempdir() |
| 134 | +lib.save(temp.relpath("graphlib.o")) |
| 135 | +remote.upload(temp.relpath("graphlib.o")) |
| 136 | +lib = remote.load_module("graphlib.o") |
| 137 | +ctx = remote.ext_dev(0) if "vta" in target else remote.cpu(0) |
| 138 | + |
| 139 | +print("Build complete...") |
| 140 | + |
| 141 | +def run_e2e(graph): |
| 142 | + """Running end to end example |
| 143 | + """ |
| 144 | + if debug_fpga_only: |
| 145 | + graph = mark_nop(graph, skip_conv_layer=(0,)) |
| 146 | + m = graph_runtime.create(graph, lib, ctx) |
| 147 | + # set inputs |
| 148 | + m.set_input('data', tvm.nd.array(x.astype("float32"))) |
| 149 | + m.set_input(**params) |
| 150 | + # execute |
| 151 | + timer = m.module.time_evaluator("run", ctx, number=10) |
| 152 | + tcost = timer() |
| 153 | + # get outputs |
| 154 | + tvm_output = m.get_output( |
| 155 | + 0,tvm.nd.empty((1000,), dtype, remote.cpu(0))) |
| 156 | + top1 = np.argmax(tvm_output.asnumpy()) |
| 157 | + print('TVM prediction top-1:', top1, synset[top1]) |
| 158 | + print("t-cost=%g" % tcost.mean) |
| 159 | + |
| 160 | + |
| 161 | +def run_layer(old_graph): |
| 162 | + """Run a certain layer.""" |
| 163 | + for layer_id in range(1, 2): |
| 164 | + graph = mark_nop(old_graph, layer_id) |
| 165 | + m = graph_runtime.create(graph, lib, ctx) |
| 166 | + # set inputs |
| 167 | + m.set_input('data', tvm.nd.array(x.astype("float32"))) |
| 168 | + m.set_input(**params) |
| 169 | + # execute |
| 170 | + timer = m.module.time_evaluator("run", ctx, number=10) |
| 171 | + tcost = timer() |
| 172 | + print("resnet[%d]: %g\n"% (layer_id, tcost.mean)) |
| 173 | + |
| 174 | +run_e2e(graph) |
0 commit comments