chainer-nnvm-example/run_pc.py at master · tkat0/chainer-nnvm-example · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import sys

import numpy as np
import onnx
import tvm
from tvm.contrib import graph_runtime, rpc
import nnvm.frontend
import nnvm.compiler

from PIL import Image
from matplotlib import pyplot as plt

from models.YOLOv2_tiny.model import postprocess

num_iter = 100
dtype = np.float32
opt_level = 3

onnx_graph = onnx.load('./models/YOLOv2_tiny/YOLOv2_tiny.onnx')

org_img = Image.open('./data/test.jpg')
org_img = org_img.resize((352, 352))
img = np.asarray(org_img).astype(np.float32).copy()
img = img.transpose(2,0,1)
img /= 255.
img = img[np.newaxis,:]

org_img = np.asarray(org_img).astype(np.float32)
org_img = org_img.transpose(2,0,1)

data_array = img
input_name = 'input_0'
data_shape = data_array.shape
out_shape = (1,125,352//32,352//32)

# GET model from frameworks
# change xyz to supported framework name.
sym, params = nnvm.frontend.from_onnx(onnx_graph)
#print(sym.debug_str())

# CPU
# OPTIMIZE and COMPILE the graph to get a deployable module
# target can be "opencl", "llvm", "metal" or any target supported by tvm
target = "llvm"
with nnvm.compiler.build_config(opt_level=opt_level, add_pass=None):
    graph, lib, params = nnvm.compiler.build(sym, target, {input_name: data_shape}, params=params)

# DEPLOY and run on cpu(0)
module = graph_runtime.create(graph, lib, tvm.cpu(0))
module.set_input('input_0', tvm.nd.array(data_array.astype(dtype)))
module.set_input(**params)
module.run()

output = tvm.nd.empty(out_shape, ctx=tvm.cpu(0))
output = module.get_output(0, output).asnumpy()
#np.save('output-nnvm-darwin-cpu.npy', output)

print('benchmark on cpu')
ctx = tvm.cpu(0)
ftimer = module.module.time_evaluator("run", ctx, num_iter)
prof_res = ftimer()
print(prof_res)

# GPU
sym, params = nnvm.frontend.from_onnx(onnx_graph)

target = "opencl"
with nnvm.compiler.build_config(opt_level=opt_level, add_pass=None):
    graph, lib, params = nnvm.compiler.build(sym, target, {input_name: data_shape}, params=params)

# DEPLOY and run on gpu(0)
ctx = tvm.context(target, 0)
module = graph_runtime.create(graph, lib, ctx)
module.set_input('input_0', tvm.nd.array(data_array.astype(dtype)))
module.set_input(**params)
module.run()

output = tvm.nd.empty(out_shape, ctx=ctx)
output = module.get_output(0, output).asnumpy()

postprocess(output[0], org_img)
plt.savefig('output-nnvm-darwin-gpu.png')

print('benchmark on gpu')
ftimer = module.module.time_evaluator("run", ctx, num_iter)
prof_res = ftimer()
print(prof_res)