def test_fold_resnet():
batch_size = 1
num_classes = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size, ) + image_shape
net, params = nnvm.testing.resnet.get_workload(batch_size=1,
image_shape=image_shape)
ishape = {"data": data_shape}
graph = nnvm.graph.create(net)
data = np.random.uniform(size=data_shape).astype("float32")
# Initial pass do shape type inference
shape, _ = graph_util.infer_shape(graph, **ishape)
ishape.update(zip(graph.index.input_names, shape))
def run_prune(graph, params, opt_level):
# Apply optimization
with nnvm.compiler.build_config(opt_level=0):
graph = nnvm.compiler.optimize(graph, ishape)
graph, params = nnvm.compiler.build_module.precompute_prune(
graph, params)
params["data"] = data
return nnvm.compiler.build_module._run_graph(graph, params)
x = run_prune(graph, params, 0)
y = run_prune(graph, params, 3)
tvm.testing.assert_allclose(y[0].asnumpy(), x[0].asnumpy())
def check_graph(sym, op_names, data_shape, params):
dtype = "float32"
targets = {"cpu": "llvm", "opencl": "opencl"}
# execute the whole graph on cpu
shape1 = {k: v for k, v in data_shape.items()}
params1 = {k: tvm.nd.array(v) for k, v in params.items()}
orig_out = execute_original_graph(sym, target="llvm", shape=shape1,
dtype=dtype, params=params1)
# annotate and compile the graph
deploy_graph, lib_dev, params = nnvm.compiler.build_heterogeneous(
sym, targets=targets, shape=data_shape, dtype=dtype, params=params,
op_names=op_names)
module = graph_runtime.create(deploy_graph, lib_dev, tvm.context("cpu"))
module.set_input(**params)
module.run()
_, oshape = graph_util.infer_shape(deploy_graph)
module_out = []
for i in range(len(sym.list_output_names())):
out = module.get_output(i, out=tvm.nd.empty(oshape[i], dtype))
module_out.append(out)
npt.assert_allclose(out.asnumpy(), orig_out[i].asnumpy(),
rtol=1e-5, atol=1e-5)
def test_fold_resnet():
batch_size = 1
num_classes = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size,) +image_shape
net, params = nnvm.testing.resnet.get_workload(
batch_size=1, image_shape=image_shape)
ishape = {"data" : data_shape}
graph = nnvm.graph.create(net)
data = np.random.uniform(size=data_shape).astype("float32")
# Initial pass do shape type inference
shape, _ = graph_util.infer_shape(graph, **ishape)
ishape.update(zip(graph.index.input_names, shape))
def run_prune(graph, params, opt_level):
# Apply optimization
with nnvm.compiler.build_config(opt_level=0):
graph = nnvm.compiler.optimize(graph, ishape)
graph, params = nnvm.compiler.build_module.precompute_prune(graph, params)
params["data"] = data
return nnvm.compiler.build_module._run_graph(graph, params)
x = run_prune(graph, params, 0)
y = run_prune(graph, params, 3)
tvm.testing.assert_allclose(y[0].asnumpy(), x[0].asnumpy())
def check_load_module():
temp = util.tempdir()
path_lib = temp.relpath("deploy.so")
libmod.export_library(path_lib)
with open(temp.relpath("deploy.json"), "w") as fo:
fo.write(deploy_graph.json())
with open(temp.relpath("deploy.params"), "wb") as fo:
fo.write(nnvm.compiler.save_param_dict(params))
# Load lib, json, and params back.
loaded_lib = tvm.module.load(path_lib)
loaded_json = open(temp.relpath("deploy.json")).read()
loaded_json = graph.load_json(loaded_json)
loaded_params = bytearray(
open(temp.relpath("deploy.params"), "rb").read())
module = graph_runtime.create(loaded_json, loaded_lib, contexts)
loaded_params = nnvm.compiler.load_param_dict(loaded_params)
module.set_input(**loaded_params)
module.run()
_, oshape = graph_util.infer_shape(loaded_json)
module_out = []
for i in range(len(sym.list_output_names())):
out = module.get_output(i, out=tvm.nd.empty(oshape[i], dtype))
module_out.append(out)
npt.assert_allclose(out.asnumpy(),
orig_out[i].asnumpy(),
rtol=1e-5,
atol=1e-5)
def test_infer_attr():
x = sym.Variable("x")
y = x * 2
g = nnvm.graph.create(y)
ishape, oshape = graph_util.infer_shape(g, x=(10, 20))
assert tuple(oshape[0]) == (10, 20)
itype, otype = graph_util.infer_dtype(g, x="float32")
assert otype[0] == "float32"
def test_infer_attr():
x = sym.Variable("x")
y = x * 2
g = nnvm.graph.create(y)
ishape, oshape = graph_util.infer_shape(g, x=(10,20))
assert tuple(oshape[0]) == (10, 20)
itype, otype = graph_util.infer_dtype(g, x="float32")
assert otype[0] == "float32"
def infer_channels(inputs, params, transpose=False):
"""A hack for getting 'channels' or 'units' since caffe2 don't provide
these attributes. We check the shape of weights provided to get the number.
"""
g = _graph.create(inputs)
shape_dict = {k: v.shape for k, v in params.items()}
_, out_shapes = graph_util.infer_shape(g, **shape_dict)
channels = out_shapes[0][0] if not transpose else out_shapes[0][1]
return channels
def check_inmemory_module():
module = graph_runtime.create(deploy_graph, libmod, contexts)
module.set_input(**params)
module.run()
_, oshape = graph_util.infer_shape(deploy_graph)
module_out = []
for i in range(len(sym.list_output_names())):
out = module.get_output(i, out=tvm.nd.empty(oshape[i], dtype))
module_out.append(out)
npt.assert_allclose(out.asnumpy(), orig_out[i].asnumpy(),
rtol=1e-5, atol=1e-5)
def init_params(net, input_name, input_shape, dtype):
params = {}
g = graph.create(net)
input_shapes, _ = graph_util.infer_shape(g, data=input_shape)
shape_dict = dict(zip(g.index.input_names, input_shapes))
for k, v in shape_dict.items():
if k == input_name:
continue
init_value = np.random.random(v).astype(dtype)
params[k] = tvm.nd.array(init_value, ctx=tvm.cpu(0))
return params
def test_multi_loss_graph_gradients():
# input data
shape1 = (1000, 100)
data1 = sym.Variable('data1', shape=(1000, 100), dtype=0)
# fake non-sparse label
label = sym.full(fill_value=3)
# square loss
sub1 = sym.elemwise_sub(data1, label, name="sub1")
square_loss = sym.sum(data=sub1**2, axis=1, name="square_loss")
# fake loss1
shape2 = (1000, )
data2 = sym.Variable('data2', shape=shape2, dtype=0)
loss1 = sym.sqrt(data2, name="loss1")
# fake loss2
loss2 = sym.relu(data1, name='loss2')
# block loss1
total_loss = sym.elemwise_sum(sym.block_grad(loss1),
square_loss,
num_args=2,
name="total_loss")
# grad_g.symbol.list_output_names()
# >> ['loss1_grad_0_output', 'grad_sum_output']
grad_g = graph_util.get_gradient_graph([total_loss, loss2],
total_loss.list_input_variables())
# infer shape
in_shapes, out_shapes = graph_util.infer_shape(grad_g)
assert out_shapes == [list(shape2), list(shape1)]
# grad_data1 is elemwise_sum of grad_loss2, grad_square_loss
grad_data1 = grad_g.symbol[1]
assert grad_data1.list_attr()['num_args'] == '2'
# block grad should return zero grad
grad_data2 = grad_g.symbol[0]
assert 'zeros_like' in grad_g.ir()
# test reverse infer shape for label
assert grad_g.apply('InferShape').json_attr('shape_num_unknown_nodes') == 0
# infer type
in_dtypes, out_dtypes = graph_util.infer_dtype(grad_g)
assert out_dtypes == ['float32', 'float32']
# test reverse infer type for label
assert grad_g.apply('InferType').json_attr('dtype_num_unknown_nodes') == 0
def test_multi_loss_graph_gradients():
# input data
shape1 = (1000, 100)
data1 = sym.Variable('data1', shape=(1000, 100), dtype=0)
# fake non-sparse label
label = sym.full(fill_value=3)
# square loss
sub1 = sym.elemwise_sub(data1, label, name="sub1")
square_loss = sym.sum(data=sub1**2, axis=1, name="square_loss")
# fake loss1
shape2 = (1000, )
data2 = sym.Variable('data2', shape=shape2, dtype=0)
loss1 = sym.sqrt(data2, name="loss1")
# fake loss2
loss2 = sym.relu(data1, name='loss2')
# block loss1
total_loss = sym.elemwise_sum(
sym.block_grad(loss1),
square_loss,
num_args=2,
name="total_loss")
# grad_g.symbol.list_output_names()
# >> ['loss1_grad_0_output', 'grad_sum_output']
grad_g = graph_util.get_gradient_graph([total_loss, loss2], total_loss.list_input_variables())
# infer shape
in_shapes, out_shapes = graph_util.infer_shape(grad_g)
assert out_shapes == [list(shape2), list(shape1)]
# grad_data1 is elemwise_sum of grad_loss2, grad_square_loss
grad_data1 = grad_g.symbol[1]
assert grad_data1.list_attr()['num_args'] == '2'
# block grad should return zero grad
grad_data2 = grad_g.symbol[0]
assert 'zeros_like' in grad_g.ir()
# test reverse infer shape for label
assert grad_g.apply('InferShape').json_attr('shape_num_unknown_nodes') == 0
# infer type
in_dtypes, out_dtypes = graph_util.infer_dtype(grad_g)
assert out_dtypes == ['float32', 'float32']
# test reverse infer type for label
assert grad_g.apply('InferType').json_attr('dtype_num_unknown_nodes') == 0
def _init_params(graph, input_shapes, initializer=init.Xavier(), seed=10):
if isinstance(graph, sym.Symbol):
graph = nnvm.graph.create(graph)
ishapes, _ = graph_util.infer_shape(graph, **input_shapes)
param_shapes = dict(zip(graph.index.input_names, ishapes))
np.random.seed(seed)
params = {}
for param, shape in param_shapes.items():
if param in {'data', 'label'} or not shape:
continue
init_value = np.empty(shape).astype('float32')
initializer(param, init_value)
params[param] = tvm.nd.array(init_value)
return params
def _init_params(graph, input_shapes, initializer=init.Xavier(), seed=10):
if isinstance(graph, sym.Symbol):
graph = nnvm.graph.create(graph)
ishapes, _ = graph_util.infer_shape(graph, **input_shapes)
param_shapes = dict(zip(graph.index.input_names, ishapes))
np.random.seed(seed)
params = {}
for param, shape in param_shapes.items():
if param in {'data', 'label'} or not shape:
continue
init_value = np.empty(shape).astype('float32')
initializer(param, init_value)
params[param] = tvm.nd.array(init_value)
return params
def execute_original_graph(sym, target, shape, dtype, params):
subgraph = graph.create(sym)
deploy_graph, lib, params = nnvm.compiler.build(
subgraph, target=target, shape=shape, dtype=dtype, params=params)
ctx = tvm.cpu()
module = graph_runtime.create(deploy_graph, lib, ctx)
module.set_input(**params)
module.run()
_, oshape = graph_util.infer_shape(deploy_graph)
module_out = []
for i in range(len(sym.list_output_names())):
out = module.get_output(i, out=tvm.nd.empty(oshape[i], dtype))
module_out.append(out)
return module_out
def heterogeneous_ssd(sym, op_names, data_shape=None, params=None,
test_image_path=None):
target, dtype = "llvm", "float32"
# targets = {"cpu": "llvm", "opencl": str(
# tvm.target.intel_graphics())}
targets = {"cpu": "llvm", "opencl": "opencl"}
with nnvm.compiler.build_config(opt_level = 3):
deploy_graph, lib_dev, params = \
nnvm.compiler.build_heterogeneous(sym, targets=targets,
shape=data_shape,
dtype=dtype, params=params,
op_names=op_names)
# import sys
# sys.stdout = open("annotated.json", "w")
# print(deploy_graph.json)
host_ctx = tvm.context("cpu")
module = graph_runtime.create(deploy_graph, lib_dev, host_ctx)
dshape = data_shape["data"]
# Preprocess image
image = cv2.imread(test_image_path)
img_data = cv2.resize(image, (dshape[2], dshape[3]))
img_data = img_data[:, :, (2, 1, 0)].astype(np.float32)
img_data -= np.array([123, 117, 104])
img_data = np.transpose(np.array(img_data), (2, 0, 1))
img_data = np.expand_dims(img_data, axis=0)
module.set_input('data', tvm.nd.array(img_data.astype(dtype)))
module.set_input(**params)
module.run()
_, oshape = graph_util.infer_shape(
deploy_graph, shape={"data": dshape})
tvm_output = module.get_output(
0, tvm.nd.empty(tuple(oshape[0]), dtype))
ftimer = module.module.time_evaluator("run", host_ctx, 2)
for i in range(2):
prof_res = ftimer()
print(prof_res)
# sleep for avoiding device overheat
if i + 1 != 5:
time.sleep(45)
return image, tvm_output
def test_cnn_gradients():
# input data
h = 128
w = 128
data_shape = (1000, 3, h, w)
data = sym.Variable('data', shape=data_shape, dtype=0)
# conv2d
num_channels = 64
kernel_size = 32
conv_w_shape = (num_channels, 3, kernel_size, kernel_size)
conv_b_shape = (num_channels,)
conv_w = sym.Variable('conv_w', shape=conv_w_shape)
conv_b = sym.Variable('conv_b', shape=conv_b_shape)
conv1 = sym.conv2d(data=data, weight=conv_w, bias=conv_b,
channels=num_channels, kernel_size=(kernel_size, kernel_size),
name='conv1')
# relu1
relu1 = sym.relu(data=conv1, name='relu1')
# max pooling
max_pooling1 = sym.max_pool2d(data=relu1, pool_size=(2, 2), name='max_pooling1')
# flatten
flatten1 = sym.flatten(data=max_pooling1)
# shape after flatten
flatten_out_shape = (h - kernel_size) * (w - kernel_size) * num_channels
# dense1
dense1_hidden_units = 100
dense1 = sym.dense(data=flatten1, name='dense1', units=dense1_hidden_units)
# relu2
relu2 = sym.relu(data=dense1, name='relu2')
# dense2
dense2_hidden_units = 10
dense2 = sym.dense(data=relu2, name='dense2', units=dense2_hidden_units)
# softmax
mlp = sym.softmax(data=dense2, name='softmax')
# fake non-sparse label
label = sym.full_like(mlp, fill_value=1)
# cross entropy loss
ce_loss = sym.sum(
sym.elemwise_mul(sym.log_softmax(dense2), label),
axis=1,
keepdims=True,
name="ce_loss")
# input variables:
# print grad_g.symbol.list_input_names()
# >> ['data', 'conv_w', 'conv_b',
# 'dense1_weight', 'dense1_bias',
# 'dense2_weight', 'dense2_bias']
# output gradient variables:
# print grad_g.symbol.list_output_names()
# >> ['conv1_grad_data', 'conv1_grad_weight', 'conv1_grad_bias',
# 'dense1_grad_weight', 'dense1_grad_bias',
# 'dense2_grad_weight', 'dense2_grad_bias']
grad_g = graph_util.get_gradient_graph(ce_loss, ce_loss.list_input_variables())
# infer shape
in_shapes, out_shapes = graph_util.infer_shape(grad_g)
# forward graph shape
assert in_shapes == [list(data_shape), list(conv_w_shape), list(conv_b_shape),
[dense1_hidden_units, flatten_out_shape], [dense1_hidden_units],
[dense2_hidden_units, dense1_hidden_units], [dense2_hidden_units]]
# input grads shape should be equal with input shape
assert in_shapes == out_shapes
# output grads w.r.t input variables
grads = graph_util.gradients(ce_loss, ce_loss.list_input_variables())
# gradients number should be equal with grad_input number
assert len(grads) == len(ce_loss.list_input_variables())
# infer type
in_dtypes, out_dtypes = graph_util.infer_dtype(grad_g)
assert out_dtypes == ['float32', 'float32', 'float32', 'float32', 'float32', 'float32', 'float32']
def _get_shape(sym, shape_dict):
"""Get the shape of a node.
"""
return graph_util.infer_shape(
nnvm.graph.create(sym), **shape_dict)[1][0]
def _get_shape(sym, shape_dict):
"""Get the shape of a node.
"""
return graph_util.infer_shape(
nnvm.graph.create(sym), **shape_dict)[1][0]
def test_cnn_gradients():
# input data
h = 128
w = 128
data_shape = (1000, 3, h, w)
data = sym.Variable('data', shape=data_shape, dtype=0)
# conv2d
num_channels = 64
kernel_size = 32
conv_w_shape = (num_channels, 3, kernel_size, kernel_size)
conv_b_shape = (num_channels, )
conv_w = sym.Variable('conv_w', shape=conv_w_shape)
conv_b = sym.Variable('conv_b', shape=conv_b_shape)
conv1 = sym.conv2d(data=data,
weight=conv_w,
bias=conv_b,
channels=num_channels,
kernel_size=(kernel_size, kernel_size),
name='conv1')
# relu1
relu1 = sym.relu(data=conv1, name='relu1')
# max pooling
max_pooling1 = sym.max_pool2d(data=relu1,
pool_size=(2, 2),
name='max_pooling1')
# flatten
flatten1 = sym.flatten(data=max_pooling1)
# shape after flatten
flatten_out_shape = (h - kernel_size) * (w - kernel_size) * num_channels
# dense1
dense1_hidden_units = 100
dense1 = sym.dense(data=flatten1, name='dense1', units=dense1_hidden_units)
# relu2
relu2 = sym.relu(data=dense1, name='relu2')
# dense2
dense2_hidden_units = 10
dense2 = sym.dense(data=relu2, name='dense2', units=dense2_hidden_units)
# softmax
mlp = sym.softmax(data=dense2, name='softmax')
# fake non-sparse label
label = sym.full_like(mlp, fill_value=1)
# cross entropy loss
ce_loss = sym.sum(sym.elemwise_mul(sym.log_softmax(dense2), label),
axis=1,
keepdims=True,
name="ce_loss")
# input variables:
# print grad_g.symbol.list_input_names()
# >> ['data', 'conv_w', 'conv_b',
# 'dense1_weight', 'dense1_bias',
# 'dense2_weight', 'dense2_bias']
# output gradient variables:
# print grad_g.symbol.list_output_names()
# >> ['conv1_grad_data', 'conv1_grad_weight', 'conv1_grad_bias',
# 'dense1_grad_weight', 'dense1_grad_bias',
# 'dense2_grad_weight', 'dense2_grad_bias']
grad_g = graph_util.get_gradient_graph(ce_loss,
ce_loss.list_input_variables())
# infer shape
in_shapes, out_shapes = graph_util.infer_shape(grad_g)
# forward graph shape
assert in_shapes == [
list(data_shape),
list(conv_w_shape),
list(conv_b_shape), [dense1_hidden_units, flatten_out_shape],
[dense1_hidden_units], [dense2_hidden_units, dense1_hidden_units],
[dense2_hidden_units]
]
# input grads shape should be equal with input shape
assert in_shapes == out_shapes
# output grads w.r.t input variables
grads = graph_util.gradients(ce_loss, ce_loss.list_input_variables())
# gradients number should be equal with grad_input number
assert len(grads) == len(ce_loss.list_input_variables())
# infer type
in_dtypes, out_dtypes = graph_util.infer_dtype(grad_g)
assert out_dtypes == [
'float32', 'float32', 'float32', 'float32', 'float32', 'float32',
'float32'
]
net = symbol.Group([net1, net2])
deploy_graph, lib, params = nnvm.compiler.build(
net, target="llvm", shape=shape_dict, dtype="float32", params=params)
temp = path.curdir
path_lib = path.join(temp, "deploy.so")
lib.export_library(path_lib)
with open(path.join(temp, "deploy.json"), "w") as fo:
fo.write(deploy_graph.json())
with open(path.join(temp, "deploy.params"), "wb") as fo:
fo.write(nnvm.compiler.save_param_dict(params))
loaded_lib = tvm.module.load(path_lib)
loaded_json = open(path.join(temp, "deploy.json")).read()
loaded_json = graph.load_json(loaded_json)
loaded_params = bytearray(
open(path.join(temp, "deploy.params"), "rb").read())
module = graph_runtime.create(loaded_json, loaded_lib, tvm.cpu(0))
loaded_params = nnvm.compiler.load_param_dict(loaded_params)
module.set_input(**loaded_params)
module.run()
_, oshape = graph_util.infer_shape(loaded_json)
out1 = module.get_output(0, out=tvm.nd.empty(oshape[0], "float32"))
assert np.allclose(data1 + data2, out1.asnumpy())
out2 = module.get_output(1, out=tvm.nd.empty(oshape[1], "float32"))
assert np.allclose(data3 + data4, out2.asnumpy())
