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])
def test_local_save_load(): if not tvm.module.enabled("opengl"): return if not tvm.module.enabled("llvm"): return n = tvm.var("n") A = tvm.placeholder((n, ), name='A', dtype='int32') B = tvm.placeholder((n, ), name='B', dtype='int32') C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C") s = tvm.create_schedule(C.op) s[C].opengl() f = tvm.build(s, [A, B, C], "opengl", target_host="llvm", name="myadd") ctx = tvm.opengl(0) n = 10 a = tvm.nd.array(np.random.uniform(high=10, size=(n)).astype(A.dtype), ctx) b = tvm.nd.array(np.random.uniform(high=10, size=(n)).astype(B.dtype), ctx) c = tvm.nd.array(np.zeros((n), dtype=C.dtype), ctx) f(a, b, c) temp = util.tempdir() path_so = temp.relpath("myadd.so") f.export_library(path_so) f1 = tvm.module.load(path_so) f1(a, b, c) np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
def test_local_save_load(): if not tvm.module.enabled("opengl"): return if not tvm.module.enabled("llvm"): return n = tvm.var("n") A = tvm.placeholder((n,), name='A', dtype='int32') B = tvm.placeholder((n,), name='B', dtype='int32') C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C") s = tvm.create_schedule(C.op) s[C].opengl() f = tvm.build(s, [A, B, C], "opengl", target_host="llvm", name="myadd") ctx = tvm.opengl(0) n = 10 a = tvm.nd.array(np.random.uniform(high=10, size=(n)).astype(A.dtype), ctx) b = tvm.nd.array(np.random.uniform(high=10, size=(n)).astype(B.dtype), ctx) c = tvm.nd.array(np.zeros((n), dtype=C.dtype), ctx) f(a, b, c) temp = util.tempdir() path_so = temp.relpath("myadd.so") f.export_library(path_so) f1 = tvm.module.load(path_so) f1(a, b, c) tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
def test_local_gemm(): if not tvm.module.enabled("opengl"): return if not tvm.module.enabled("llvm"): return nn = 1024 n = tvm.var('n') n = tvm.convert(nn) m = n l = n A = tvm.placeholder((n, l), name='A', dtype='int32') B = tvm.placeholder((m, l), name='B', dtype='int32') k = tvm.reduce_axis((0, l), name='k') C = tvm.compute((n, m), lambda ii, jj: tvm.sum(A[ii, k] * B[jj, k], axis=k), name='CC') s = tvm.create_schedule(C.op) s[C].opengl() print(tvm.lower(s, [A, B, C], simple_mode=True)) f = tvm.build(s, [A, B, C], "opengl", name="gemm") print("------opengl code------") print(f.imported_modules[0].get_source(fmt="gl")) ctx = tvm.opengl() n, m, l = nn, nn, nn a_np = np.random.uniform(low=0, high=10, size=(n, l)).astype(A.dtype) b_np = np.random.uniform(low=0, high=10, size=(m, l)).astype(B.dtype) a = tvm.nd.array(a_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros((n, m), dtype=C.dtype), ctx) f(a, b, c) tvm.testing.assert_allclose(c.asnumpy(), np.dot(a_np, b_np.T))
def test_local_gemm(): if not tvm.runtime.enabled("opengl"): return if not tvm.runtime.enabled("llvm"): return nn = 1024 n = te.var('n') n = tvm.runtime.convert(nn) m = n l = n A = te.placeholder((n, l), name='A', dtype='int32') B = te.placeholder((m, l), name='B', dtype='int32') k = te.reduce_axis((0, l), name='k') C = te.compute((n, m), lambda ii, jj: te.sum(A[ii, k] * B[jj, k], axis=k), name='CC') s = te.create_schedule(C.op) s[C].opengl() print(tvm.lower(s, [A, B, C], simple_mode=True)) f = tvm.build(s, [A, B, C], "opengl", name="gemm") print("------opengl code------") print(f.imported_modules[0].get_source(fmt="gl")) ctx = tvm.opengl() n, m, l = nn, nn, nn a_np = np.random.uniform(low=0, high=10, size=(n, l)).astype(A.dtype) b_np = np.random.uniform(low=0, high=10, size=(m, l)).astype(B.dtype) a = tvm.nd.array(a_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros((n, m), dtype=C.dtype), ctx) f(a, b, c) tvm.testing.assert_allclose(c.asnumpy(), np.dot(a_np, b_np.T))
def test_local_multi_stage(): if not tvm.module.enabled("opengl"): return if not tvm.module.enabled("llvm"): return n = tvm.var("n") A = tvm.placeholder((n,), name='A', dtype="int32") B = tvm.compute((n,), lambda i: A[i] + 1, name="B") C = tvm.compute((n,), lambda i: B[i] * 2, name="C") s = tvm.create_schedule(C.op) s[B].opengl() s[C].opengl() f = tvm.build(s, [A, C], "opengl", name="multi_stage") ctx = tvm.opengl(0) n = 10 a = tvm.nd.array(np.random.uniform(size=(n,)).astype(A.dtype), ctx) c = tvm.nd.array(np.random.uniform(size=(n,)).astype(B.dtype), ctx) f(a, c) tvm.testing.assert_allclose(c.asnumpy(), (a.asnumpy() + 1) * 2)
def test_local_multi_stage(): if not tvm.module.enabled("opengl"): return if not tvm.module.enabled("llvm"): return n = tvm.var("n") A = tvm.placeholder((n, ), name='A', dtype="int32") B = tvm.compute((n, ), lambda i: A[i] + 1, name="B") C = tvm.compute((n, ), lambda i: B[i] * 2, name="C") s = tvm.create_schedule(C.op) s[B].opengl() s[C].opengl() f = tvm.build(s, [A, C], "opengl", name="multi_stage") ctx = tvm.opengl(0) n = 10 a = tvm.nd.array(np.random.uniform(size=(n, )).astype(A.dtype), ctx) c = tvm.nd.array(np.random.uniform(size=(n, )).astype(B.dtype), ctx) f(a, c) tvm.testing.assert_allclose(c.asnumpy(), (a.asnumpy() + 1) * 2)
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])
exit() else: lib.export_library('from_mxnet.so') ###################################################################### # Execute the portable graph on TVM # --------------------------------- # Now, we would like to reproduce the same forward computation using TVM. from tvm.contrib import graph_runtime if target == 'llvm': ctx = tvm.cpu(0) elif target == 'cuda': ctx = tvm.gpu(0) elif target == 'opengl': ctx = tvm.opengl(0) elif target == 'opencl': ctx = tvm.cl(0) elif target == 'vulkan': ctx = tvm.vulkan(0) elif target == 'metal': ctx = tvm.metal(0) else: raise ValueError('No supported context type for ' % target) dtype = 'float32' m = graph_runtime.create(graph, lib, ctx) # set inputs m.set_input('data', tvm.nd.array(x.astype(dtype))) m.set_input(**params) # execute