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())