README.md 1.08 KB
Newer Older
tqchen committed
1
# NNVM Compiler Module of TVM Stack
2 3 4 5 6 7 8

```python
import tvm
from tvm.contrib import graph_runtime, rpc
import nnvm.frontend
import nnvm.compiler

9
# GET model from frameworks
10 11 12
# change xyz to supported framework name.
graph, params = nnvm.frontend.from_xyz(...)

13
# OPTIMIZE and COMPILE the graph to get a deployable module
14 15
# target can be "opencl", "llvm", "metal" or any target supported by tvm
target = "cuda"
16
graph, lib, params = nnvm.compiler.build(graph, target, {"data", data_shape}, params=params)
17

18
# DEPLOY and run on gpu(0)
19 20
module = graph_runtime.create(graph, lib, tvm.gpu(0))
module.set_input(**params)
21
module.run(data=data_array)
22
output = tvm.nd.empty(out_shape, ctx=tvm.gpu(0))
23
module.get_output(0, output)
24

25
# DEPLOY to REMOTE mobile/rasp/browser with minimum tvm rpc runtime
26 27 28 29 30 31 32 33 34 35
# useful for quick experiments on mobile devices
remote = rpc.connect(remote_host, remote_port)
lib.export_library("mylib.so")
remote.upload("mylib.so")
rlib = rpc.load_module("mylib.so")
# run on remote device
rmodule = graph_runtime.create(graph, rlib, remote.gpu(0))
rmodule.set_input(**params)
rmodule.run()
```