def compile_run_graph(device, target): if not tvm.module.enabled(device): print("Skip test because %s is not enabled." % device) return out_channels = 16 data1 = symbol.Variable(name="data1") data2 = symbol.Variable(name="data2") simple_net1 = symbol.conv2d(data=data1, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) simple_net2 = symbol.conv2d(data=data2, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) ret = symbol.elemwise_add(simple_net1, simple_net2) ret = symbol.conv2d(ret, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) batch_size = 1 data_shape = (batch_size, 3, 224, 224) shape_dict = {"data1": data_shape, "data2": data_shape} params = {} params["data1"] = np.random.uniform(-1, 1, size=data_shape).astype("float32") params["data2"] = np.random.uniform(-1, 1, size=data_shape).astype("float32") op_name_device = {"elemwise_add": "cpu", "conv2d": device} fallback_device = tvm.context("cpu") target = {"cpu": "llvm", device: target} # No op will be fused. 3 additional device copy nodes are required. check_annotated_graph(ret, target, op_name_device, 15, fallback_device, shape_dict, params)
def sample(): x = sym.Variable("x") y = sym.Variable("y") z1 = sym.elemwise_add(x, sym.sqrt(y)) z2 = sym.log(x) gradient = graph_util.gradients([z1, z2], [x, y]) print(gradient)
def test_concatenate_conv2d(): ch = 3 size = 8 data = sym.Variable(name="data") concat = sym.concatenate(data, data, axis=1) conv = sym.conv2d(data=concat, kernel_size=(1,1), channels=ch*2, use_bias=False, name="conv") net = sym.elemwise_add(concat, conv) dtype="float32" dshape = (1, ch, size, size) kshape = (ch*2, ch*2, 1, 1) oshape = (1, ch*2, size, size) shape_dict = {"data": dshape} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(net, target, shape_dict) # data, conv weight, conv op, concat assert graph.index.num_nodes == 4 data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) kernel = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) m = graph_runtime.create(graph, lib, ctx) m.run(data=data, conv_weight=kernel) # get output out = m.get_output(0, tvm.nd.empty(oshape, dtype)) concat = np.concatenate((data.asnumpy(), data.asnumpy()), axis=1) conv = topi.testing.conv2d_nchw_python( concat, kernel.asnumpy(), (1,1), 'SAME') ref = concat + conv tvm.testing.assert_allclose(out.asnumpy(), ref, rtol=1e-5)
def test_conv_network(): """ The network is as following: data1 data2 | | conv2d conv2d \ / elemwise_add | conv2d """ out_channels = 16 data1 = symbol.Variable(name="data1") data2 = symbol.Variable(name="data2") simple_net1 = symbol.conv2d(data=data1, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) simple_net2 = symbol.conv2d(data=data2, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) ret = symbol.elemwise_add(simple_net1, simple_net2) ret = symbol.conv2d(ret, kernel_size=(3, 3), channels=out_channels, padding=(1, 1), use_bias=True) batch_size = 1 data_shape = (batch_size, 3, 224, 224) shape_dict = {"data1": data_shape, "data2": data_shape} params = {} params["data1"] = np.random.uniform(-1, 1, size=data_shape).astype("float32") params["data2"] = np.random.uniform(-1, 1, size=data_shape).astype("float32") # No op will be fused. 3 additional device copy nodes are required. check_annotated_graph(ret, ["elemwise_add"], 15, shape_dict, params)
def test_concatenate_conv2d(): ch = 3 size = 8 data = sym.Variable(name="data") concat = sym.concatenate(data, data, axis=1) conv = sym.conv2d(data=concat, kernel_size=(1, 1), channels=ch * 2, use_bias=False, name="conv") net = sym.elemwise_add(concat, conv) dtype = "float32" dshape = (1, ch, size, size) kshape = (ch * 2, ch * 2, 1, 1) oshape = (1, ch * 2, size, size) shape_dict = {"data": dshape} for target, ctx in ctx_list(): graph, lib, _ = nnvm.compiler.build(net, target, shape_dict) # data, conv weight, conv op, concat assert graph.index.num_nodes == 4 data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) kernel = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) m = graph_runtime.create(graph, lib, ctx) m.run(data=data, conv_weight=kernel) # get output out = m.get_output(0, tvm.nd.empty(oshape, dtype)) concat = np.concatenate((data.asnumpy(), data.asnumpy()), axis=1) conv = topi.testing.conv2d_nchw_python(concat, kernel.asnumpy(), (1, 1), 'SAME') ref = concat + conv np.testing.assert_allclose(out.asnumpy(), ref, rtol=1e-5)
def test_copy(): x = sym.Variable('x') z = sym.Variable('z') y = sym.exp(sym.elemwise_add(x, x, name='add', gpu=2), name='exp', gpu=1, attr={"kk": "1"}) assert y.__copy__().debug_str() == y.debug_str()
def test_plan_memory(): x = sym.Variable('x', shape=(4, 2)) x2 = sym.elemwise_add(x, x, name='addk') y = sym.flatten(x2, name="reshapek") y = sym.elemwise_add(y, x2, name="add2") y = sym.elemwise_add(y, y) g = graph.create(y) g._set_json_attr("shape_attr_key", "shape") g = g.apply(["InferShape", "InferType", "PlanMemory"]) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] storage_id = g.json_attr('storage_id') nindex = {n['name']: i for i, n in enumerate(jnodes)} assert (storage_id[jnode_row_ptr[nindex["addk"]]] != storage_id[jnode_row_ptr[nindex["reshapek"]]]) assert (storage_id[jnode_row_ptr[nindex["add2"]]] == storage_id[ jnode_row_ptr[nindex["reshapek"]]])
def test_plan_memory(): x = sym.Variable('x', shape=(4, 2)) x2 = sym.elemwise_add(x, x, name='addk') y = sym.flatten(x2, name="reshapek") y = sym.elemwise_add(y, x2, name="add2") y = sym.elemwise_add(y, y) g = graph.create(y) g._set_json_attr("shape_attr_key", "shape") g = g.apply(["InferShape", "InferType", "PlanMemory"]) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] storage_id = g.json_attr('storage_id') nindex = {n['name']: i for i, n in enumerate(jnodes)} assert (storage_id[jnode_row_ptr[nindex["addk"]]] != storage_id[jnode_row_ptr[nindex["reshapek"]]]) assert (storage_id[jnode_row_ptr[nindex["add2"]]] == storage_id[jnode_row_ptr[nindex["reshapek"]]])
def test_default_input(): x = sym.Variable('x') y = sym.dense(data=x, units=30, name='fc', use_bias=False) assert y.list_input_names() == ['x', 'fc_weight'] tname = [z.list_output_names()[0] for z in y.list_input_variables()] assert tname == y.list_input_names() try: z = sym.elemwise_add(x) assert False except NNVMError: pass
def test_compose(): x = sym.Variable('x') z = sym.Variable('z') y = sym.exp(sym.elemwise_add(x, x, name='add', gpu=2), name='exp', gpu=1, attr={"kk": "1"}) assert y.list_input_names() == ['x'] assert y.list_output_names() == ["exp_output"] assert y.list_attr()['gpu'] == '1' z = y.get_internals() assert z['add_output'].list_output_names() == ['add_output'] assert y.list_attr(recursive=True)['add$gpu'] == '2'
def simple_bn(x, gamma, beta, moving_mean, moving_var, axis=1, epsilon=1e-5, shape=None): # expect = (x - moving_mean) / sym.sqrt(moving_var + eps) * gamma + beta scale = sym.elemwise_mul(1 / sym.sqrt(moving_var + epsilon), gamma) shift = sym.elemwise_add( sym.elemwise_mul(sym.negative(moving_mean), scale), beta) # for 2D num_newaxis=len(shape) - axis - 1 if num_newaxis: scale = sym.expand_dims(scale, axis=1, num_newaxis=num_newaxis) shift = sym.expand_dims(shift, axis=1, num_newaxis=num_newaxis) return x * scale + shift
def test_infer_type(): x = sym.Variable('x', dtype=0) y = sym.elemwise_add(x, x, name='add1') y = sym.cast(y, dtype="float64", name="cast1") g = graph.create(y) g._set_json_attr("dtype_attr_key", "dtype") g = g.apply('InferType') jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('dtype')[jnode_row_ptr[nindex["cast1"]]] == 1 assert g.json_attr('dtype')[jnode_row_ptr[nindex["add1"]]] == 0
def test_infer_shape(): x = sym.Variable('x', shape=(2, 4, 2)) y = sym.elemwise_add(x, x, name='add1') y = sym.flatten(y, name="flatten") g = graph.create(y) g._set_json_attr("shape_attr_key", "shape") g = g.apply('InferShape') jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('shape')[jnode_row_ptr[nindex["flatten"]]] == [2, 8] assert g.json_attr('shape')[jnode_row_ptr[nindex["add1"]]] == [2, 4, 2]
def test_residual_block_layout_transform(): ch = 16 size = 32 data = sym.Variable(name="data") conv1 = sym.conv2d(data=data, kernel_size=(3, 3), channels=ch, padding=(1, 1), use_bias=False, name="conv1") layout_transform1 = sym.__layout_transform__(data=conv1, src_layout="NCHW", dst_layout="NCHW8c") layout_transform2 = sym.__layout_transform__(data=layout_transform1, src_layout="NCHW8c", dst_layout="NCHW") conv2 = sym.conv2d(data=conv1, kernel_size=(3, 3), channels=ch, padding=(1, 1), use_bias=False, name="conv2") elemwise_sum = sym.elemwise_add(layout_transform2, conv2) out = sym.relu(elemwise_sum) dtype = "float32" dshape = (1, ch, size, size) kshape = (ch, ch, 3, 3) oshape = (1, ch, size, size) shape_dict = {"data": dshape} target = "llvm" # only test on llvm since it involves NCHW8c layout ctx = tvm.context(target, 0) graph, lib, _ = nnvm.compiler.build(out, target, shape_dict) # data, conv1 weight, conv1, layout transform + elemwise add + relu, conv2 weight, conv2 op assert graph.index.num_nodes == 6 data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) kernel1 = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) kernel2 = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) m = graph_runtime.create(graph, lib, ctx) m.run(data=data, conv1_weight=kernel1, conv2_weight=kernel2) out = m.get_output(0, tvm.nd.empty(oshape, dtype)) conv1 = topi.testing.conv2d_nchw_python(data.asnumpy(), kernel1.asnumpy(), (1, 1), 'SAME') conv2 = topi.testing.conv2d_nchw_python(conv1, kernel2.asnumpy(), (1, 1), 'SAME') ref = np.maximum(conv1 + conv2, 0) np.testing.assert_allclose(out.asnumpy(), ref, rtol=1e-5)
def test_infer_shape_known_partial(): x = sym.Variable('x') y = sym.elemwise_add(x, x, name='add1') y = sym.flatten(y, name="flatten1") g = graph.create(y) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) shape = [[2, 4, 2], [], []] g._set_json_attr("shape", shape, 'list_shape') g = g.apply("InferShape") jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('shape')[jnode_row_ptr[nindex["flatten1"]]] == [2, 8] assert g.json_attr('shape')[jnode_row_ptr[nindex["add1"]]] == [2, 4, 2]
def test_infer_shape_known_partial(): x = sym.Variable('x') y = sym.elemwise_add(x, x, name='add1') y = sym.flatten(y, name="flatten1") g = graph.create(y) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) shape = [[2, 4, 2], [] , []] g._set_json_attr("shape", shape, 'list_shape') g = g.apply("InferShape") jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('shape')[jnode_row_ptr[nindex["flatten1"]]] == [2, 8] assert g.json_attr('shape')[jnode_row_ptr[nindex["add1"]]] == [2, 4, 2]
def test_gradient(): x = sym.Variable("x") y = sym.Variable("y") z1 = sym.elemwise_add(x, sym.sqrt(y)) z2 = sym.log(x) gradient = graph_util.gradients([z1, z2], [x, y]) assert len(gradient) == 2 g1 = sym.Variable("g1") g2 = sym.Variable("g2") grad_ys = [g1, g2] gradient = graph_util.gradients(sym.Group([z1, z2]), sym.Group([x, y]), grad_ys=grad_ys) g_graph = graph.create(sym.Group(gradient)).ir() assert len(gradient) == 2 assert "g1" in g_graph assert "g2" in g_graph
def simple_bn(x, gamma, beta, moving_mean, moving_var, axis=1, epsilon=1e-5, shape=None): # expect = (x - moving_mean) / sym.sqrt(moving_var + eps) * gamma + beta scale = sym.elemwise_mul(1 / sym.sqrt(moving_var + epsilon), gamma) shift = sym.elemwise_add( sym.elemwise_mul(sym.negative(moving_mean), scale), beta) # for 2D num_newaxis = len(shape) - axis - 1 if num_newaxis: scale = sym.expand_dims(scale, axis=1, num_newaxis=num_newaxis) shift = sym.expand_dims(shift, axis=1, num_newaxis=num_newaxis) return x * scale + shift
def test_create_full_graph(): x = sym.Variable("x") y = sym.Variable("y") z1 = sym.elemwise_add(x, sym.sqrt(y)) z2 = sym.log(x) symbol = sym.Group([z1, z2]) compute_graph = graph.create(symbol, need_backward=True) assert (compute_graph.index.num_nodes == 11) head_grads = [sym.Variable("g1"), sym.Variable("g2")] compute_graph = graph.create(symbol, need_backward=True, head_grads=head_grads) ir = compute_graph.ir() assert (compute_graph.index.num_nodes == 11) assert ("g1" in ir) assert ("g2" in ir) fixed_args = ["x"] compute_graph = graph.create(symbol, need_backward=True, fixed_args=fixed_args) assert (compute_graph.index.num_nodes == 8)
def test_residual_block_layout_transform(): ch = 16 size = 32 data = sym.Variable(name="data") conv1 = sym.conv2d(data=data, kernel_size=(3,3), channels=ch, padding = (1, 1), use_bias=False, name="conv1") layout_transform1 = sym.__layout_transform__(data=conv1, src_layout="NCHW", dst_layout="NCHW8c") layout_transform2 = sym.__layout_transform__(data=layout_transform1, src_layout="NCHW8c", dst_layout="NCHW") conv2 = sym.conv2d(data=conv1, kernel_size=(3,3), channels=ch, padding = (1, 1), use_bias=False, name="conv2") elemwise_sum = sym.elemwise_add(layout_transform2, conv2) out = sym.relu(elemwise_sum) dtype="float32" dshape = (1, ch, size, size) kshape = (ch, ch, 3, 3) oshape = (1, ch, size, size) shape_dict = {"data": dshape} target = "llvm" # only test on llvm since it involves NCHW8c layout ctx = tvm.context(target, 0) graph, lib, _ = nnvm.compiler.build(out, target, shape_dict) # data, conv1 weight, conv1, layout transform + elemwise add + relu, conv2 weight, conv2 op assert graph.index.num_nodes == 6 data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) kernel1 = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) kernel2 = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) m = graph_runtime.create(graph, lib, ctx) m.run(data=data, conv1_weight=kernel1, conv2_weight=kernel2) out = m.get_output(0, tvm.nd.empty(oshape, dtype)) conv1 = topi.testing.conv2d_nchw_python( data.asnumpy(), kernel1.asnumpy(), (1,1), 'SAME') conv2 = topi.testing.conv2d_nchw_python( conv1, kernel2.asnumpy(), (1,1), 'SAME') ref = np.maximum(conv1 + conv2, 0) tvm.testing.assert_allclose(out.asnumpy(), ref, rtol=1e-5)
# Most deep learning frameworks use computation graph to describe # their computation. In this example, we directly use # NNVM's API to construct the computational graph. # # .. note:: # # In a typical deep learning compilation workflow, # we can get the models from :any:`nnvm.frontend` # # The following code snippet describes :math:`z = x + \sqrt{y}` # and creates a nnvm graph from the description. # We can print out the graph ir to check the graph content. x = sym.Variable("x") y = sym.Variable("y") z = sym.elemwise_add(x, sym.sqrt(y)) compute_graph = nnvm.graph.create(z) print("-------compute graph-------") print(compute_graph.ir()) ###################################################################### # Compile # ------- # We can call :any:`nnvm.compiler.build` to compile the graph. # The build function takes a shape parameter which specifies the # input shape requirement. Here we only need to pass in shape of ``x`` # and the other one will be inferred automatically by NNVM. # # The function returns three values. ``deploy_graph`` contains # the final compiled graph structure. ``lib`` is a :any:`tvm.module.Module` # that contains compiled CUDA functions. We do not need the ``params``
def test_list_args(): x = sym.Variable('x') z = sym.Variable('z') y = sym.dense(data=x, name='fc', units=30) y = sym.elemwise_add(y, z, name='add1')
import nnvm.compiler import nnvm.symbol as sym import numpy as np x = sym.Variable("x") y = sym.Variable("y") z = sym.elemwise_add(x, sym.sqrt(y)) compute_graph = nnvm.graph.create(z) x_np = np.array([1, 2, 3, 4]).astype("float32") y_np = np.array([4, 4, 4, 4]).astype("float32") shape = (4, ) deploy_graph, lib, params = nnvm.compiler.build(compute_graph, target="cuda", target_host= "llvm -target=aarch64-linux-gnu", shape={"x": shape}, params={"y": y_np}, dtype="float32") x_np.tofile("./data/x.bin") np.save("./data/x.nparray", x_np) with open("./model/jetson_gpu.json", "w") as fo: fo.write(deploy_graph.json()) with open("./model/jetson_gpu.params", "wb") as fo: fo.write(nnvm.compiler.save_param_dict(params)) lib.save("./model/jetson_gpu.o") dev_modules = lib.imported_modules dev_modules[0].save("./model/jetson_gpu.ptx")