def test_alter_conv2d_layout(): data = sym.Variable("data", shape=(1, 32, 512, 512)) conv = sym.conv2d(data, name="conv", channels=16, kernel_size=(3, 3), padding=(1, 1), use_bias=False, layout="NCHW") relu = sym.relu(conv, name="relu") flatten = sym.flatten(relu, name="flatten") softmax = sym.softmax(flatten, name="softmax") g = graph.create(softmax) g = g.apply("CorrectLayout") g = graph_attr.set_dtype_inputs(g, "float32") g = g.apply(["InferShape", "InferType"]) layouts_origin = get_layouts(g) @reg.register_alter_op_layout("conv2d") def alter_conv2d_layout(attrs, inputs, tinfos): new_attrs = {k: attrs[k] for k in attrs.keys()} new_attrs["layout"] = "NCHW16c" new_attrs["kernel_layout"] = "NCHW16c" new_attrs["name"] = "conv_alter" return sym.conv2d(inputs[0], inputs[1], **new_attrs) g = g.apply("AlterOpLayout") layouts = get_layouts(g) # check copy layouts for node in ["data", "relu", "flatten", "softmax", "conv_weight"]: assert (layouts[node] == layouts_origin[node]) assert (layouts["conv_alter"] == layouts_origin["conv"])
def test_consecutive_alter_layout(): data = sym.Variable("data", shape=(1, 32, 512, 512)) pool1 = sym.global_avg_pool2d(data, name="global_avg_pool2d_1", layout="NCHW") pool2 = sym.global_avg_pool2d(pool1, name="global_avg_pool2d_2", layout="NCHW") relu = sym.relu(pool2, name="relu") g = graph.create(relu) g = g.apply("CorrectLayout") g = graph_attr.set_dtype_inputs(g, "float32") g = g.apply(["InferShape", "InferType"]) assert g.json_attr("layout") == ['NCHW', 'NCHW', 'NCHW', 'NCHW'] @reg.register_alter_op_layout("global_avg_pool2d", level=100) def alter_global_avg_pool2d_layout(attrs, inputs, tinfos): new_attrs = {k: attrs[k] for k in attrs.keys()} new_attrs["layout"] = "NCHW16c" return sym.global_avg_pool2d(inputs[0], **new_attrs) g = g.apply("AlterOpLayout") # pool1 get replaced - output layout of pool1 is not recorded # pool2 get replaced - input layout of pool2 is not recorded # thus the second entry must be undefined - it can neither recover from pool1's output, # nor from pool2's input. assert g.json_attr("layout") == ['NCHW', '__undef__', 'NCHW', 'NCHW']
def test_json_pass(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2, 2)) g = graph.create(y) ret = g.apply('SaveJSON') ret._set_json_attr('json', ret.json_attr('json')) g2 = ret.apply('LoadJSON') assert g2.apply('SaveJSON').json_attr('json') == ret.json_attr('json')
def test_json_pass(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2,2)) g = graph.create(y) ret = g.apply('SaveJSON') ret._set_json_attr('json', ret.json_attr('json')) g2 = ret.apply('LoadJSON') assert g2.apply('SaveJSON').json_attr('json') == ret.json_attr('json')
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 test_print_graph_ir(): x = sym.Variable("x", shape=(1, 1, 10, 20)) y = sym.conv2d(x + 1, name="y", channels=10, kernel_size=(3, 3)) g = graph.create(y) g = g.apply("InferShape") ir1 = g.ir() ir2 = g.ir(join_entry_attrs=["shape"]) assert ("y_bias" in ir1) assert ("shape=" in ir2)
def test_print_graph_ir(): x = sym.Variable("x", shape=(1, 1, 10, 20)) y = sym.conv2d(x + 1, name="y", channels=10, kernel_size=(3,3)) g = graph.create(y) g = g.apply("InferShape") ir1 = g.ir() ir2 = g.ir(join_entry_attrs=["shape"]) assert("y_bias" in ir1) assert("shape=" in ir2)
def test_json_pass_with_attr(): x = sym.Variable('x') y = sym.dense(data=x, name='fc', units=30) g = graph.create(y) g._set_json_attr('version', '0.1.0') ret = g.apply('SaveJSON') json_str = ret.json_attr('json') ret._set_json_attr('json', json_str) g2 = ret.apply('LoadJSON') assert g2.json_attr('version') == '0.1.0'
def test_json_pass_with_attr(): x = sym.Variable('x') y = sym.dense(data=x, name='fc', units=30) g = graph.create(y) g._set_json_attr('version', '0.1.0') ret = g.apply('SaveJSON') json_str = ret.json_attr('json') ret._set_json_attr('json', json_str) g2 = ret.apply('LoadJSON') assert g2.json_attr('version') == '0.1.0'
def from_nnvm(model, frontend, filename, input_shape=None, batch_size=None, layout="NCHW", dtype=hcl.Float(), target=None): if frontend == "keras": nnvm_model = keras.models.load_model(model) graph, params = nnvm.frontend.from_keras(nnvm_model) else: raise NameError('frontend {} is not a valid frontend'.format(frontend)) #generate json from model graph = _graph.create(graph) _json = json.loads(graph.json()) nodes = _json["nodes"] param_shape = {} #convert params to dict of numpy arrays for param in params: params[param] = params[param].asnumpy() param_shape[param] = params[param].shape #add shapes to _json for lower function _json["param_shape"] = param_shape var_sym = {} hcl.init(dtype) #generate placeholders for node in _json["arg_nodes"]: new_node = nodes[node] node_name = new_node['name'] if 'input' in node_name: gen_placeholders(var_sym, node_name, batch_size, input_shape, False) else: gen_placeholders(var_sym, node_name, batch_size, params[node_name].shape, True) args = [] for j in var_sym: args.append(var_sym[j]) func = gen_function(var_sym, nodes, _json['heads'][0][0], layout) s = gen_schedule(args, func) #transform params so they can be used in function param = [] for i in params: param.append(hcl.asarray(params[i])) if target == None: return hcl.build(s), tuple(param) else: try: f = hcl.build(s, target=target, name=filename) save_file(filename, f, target) return f, tuple(param) except ValueError: print( "target {} provided is not a compatible target".format(target))
def test_json_pass(): x = sym.Variable('x') y = sym.dense(data=x, name='conv', units=30) g = graph.create(y) ret = g.apply('SaveJSON') ret._set_json_attr('json', ret.json_attr('json')) g2 = ret.apply('LoadJSON') assert g2.apply('SaveJSON').json_attr('json') == ret.json_attr('json') json = g.json() g2 = graph.load_json(json) assert json == g2.json()
def infer_shape(sym): g = graph.create(sym) g._set_json_attr("shape_attr_key", "shape") g = g.apply("InferShape") sdict = {} vshape = g.json_attr("shape") entry_ptr = g.index.entry_ptr for i, n in enumerate(g.index.nodes): begin, end = entry_ptr[i], entry_ptr[i + 1] sdict[n["name"]] = vshape[begin:end] return sdict
def infer_shape(sym): g = graph.create(sym) g._set_json_attr("shape_attr_key", "shape") g = g.apply("InferShape") sdict = {} vshape = g.json_attr("shape") entry_ptr = g.index.entry_ptr for i, n in enumerate(g.index.nodes): begin, end = entry_ptr[i], entry_ptr[i + 1] sdict[n["name"]] = vshape[begin:end] return sdict
def test_json_pass(): x = sym.Variable('x') y = sym.dense(data=x, name='conv', units=30) g = graph.create(y) ret = g.apply('SaveJSON') ret._set_json_attr('json', ret.json_attr('json')) g2 = ret.apply('LoadJSON') assert g2.apply('SaveJSON').json_attr('json') == ret.json_attr('json') json = g.json() g2 = graph.load_json(json) assert json == g2.json()
def test_json_pass_with_attr(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2,2)) g = graph.create(y) g._set_json_attr('version', '0.1.0') ret = g.apply('SaveJSON') json_str = ret.json_attr('json') print(json_str) ret._set_json_attr('json', json_str) g2 = ret.apply('LoadJSON') assert g2.json_attr('version') == '0.1.0'
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_json_pass_with_attr(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2, 2)) g = graph.create(y) g._set_json_attr('version', '0.1.0') ret = g.apply('SaveJSON') json_str = ret.json_attr('json') print(json_str) ret._set_json_attr('json', json_str) g2 = ret.apply('LoadJSON') assert g2.json_attr('version') == '0.1.0'
def test_infer_type(): x = sym.Variable('x') y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") g = graph.create(y) 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 correct_layout(g, layout=None): if isinstance(g, nnvm.symbol.Symbol): g = graph.create(g) if layout: graph_attr.set_layout_inputs(g, layout) g = g.apply("CorrectLayout") ldict = {} vlayout = g.json_attr("layout") entry_ptr = g.index.entry_ptr for i, n in enumerate(g.index.nodes): begin, end = entry_ptr[i], entry_ptr[i + 1] ldict[n["name"]] = vlayout[begin:end] return g, ldict
def test_infer_type(): x = sym.Variable('x', dtype=0) y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, 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=(4, 2)) y = sym.add(x, x, name='add1') y = sym.reshape(y, target=(2, 4), name="reshape1") 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["reshape1"]]] == [2, 4] assert g.json_attr('shape')[jnode_row_ptr[nindex["add1"]]] == [4, 2]
def correct_layout(g, layout=None): if isinstance(g, nnvm.symbol.Symbol): g = graph.create(g) if layout: graph_attr.set_layout_inputs(g, layout) g = g.apply("CorrectLayout") ldict = {} vlayout = g.json_attr("layout") entry_ptr = g.index.entry_ptr for i, n in enumerate(g.index.nodes): begin, end = entry_ptr[i], entry_ptr[i + 1] ldict[n["name"]] = vlayout[begin:end] return g, ldict
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 gradients(ys, xs, grad_ys=None): if isinstance(ys, list): ys = symbol.Group(ys) g = graph.create(ys) g._set_symbol_list_attr('grad_ys', ys) g._set_symbol_list_attr('grad_xs', xs) ny = len(ys.list_output_names()) if grad_ys is None: grad_ys = [symbol.ones_like(ys[i]) for i in range(ny)] g._set_symbol_list_attr('grad_ys_out_grad', grad_ys) sym = g.apply('Gradient').symbol nx = len(xs) if isinstance(xs, list) else len(xs.list_output_names()) ret = [sym[i] for i in range(nx)] return ret
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_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 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 test_order_mutation_pass(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', dev='gpu') y = sym.add(y, x, name='add1') # write after read z = sym.assign(x, y, name='assign') # read after write t = sym.add(y, x, name='add2') g = graph.create(sym.Group([t, z])) jgraph = json.loads(g.apply(['OrderMutation', 'SaveJSON']).json_attr('json')) jnodes = jgraph['nodes'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert nindex['assign'] in jnodes[nindex['add2']]['control_deps'] assert nindex['conv'] in jnodes[nindex['assign']]['control_deps'] assert nindex['add1'] in jnodes[nindex['assign']]['control_deps'] assert jnodes[nindex['assign']]['inputs'][0][2] == 1
def test_order_mutation_pass(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', dev='gpu') y = sym.add(y, x, name='add1') # write after read z = sym.assign(x, y, name='assign') # read after write t = sym.add(y, x, name='add2') g = graph.create(sym.Group([t, z])) jgraph = json.loads(g.apply(['OrderMutation', 'SaveJSON']).json_attr('json')) jnodes = jgraph['nodes'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert nindex['assign'] in jnodes[nindex['add2']]['control_deps'] assert nindex['conv'] in jnodes[nindex['assign']]['control_deps'] assert nindex['add1'] in jnodes[nindex['assign']]['control_deps'] assert jnodes[nindex['assign']]['inputs'][0][2] == 1
def build(sym, target, shape, dtype="float32"): # Do shape inference in python. g = graph.create(sym) 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)} list_shape = [[]] * jnode_row_ptr[-1] list_dtype = [DTYPE_DICT[dtype]] * jnode_row_ptr[-1] for k, v in shape.items(): list_shape[jnode_row_ptr[nindex[k]]] = v g._set_json_attr("shape", list_shape, 'list_shape') g._set_json_attr("dtype", list_dtype, 'list_int') g._set_json_attr("target", target, 'str') g = g.apply("InferShape").apply("InferType") g = g.apply("GraphPartition").apply("GraphFuse") return g
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 test_place_device(): x = sym.Variable('x', device_group="stage1") y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") z = sym.add(y, y, device_group="stage2", name="add2") z = sym.add(z, sym.exp(y, device_group="stage2"), name="add3") g = graph.create(z) g._set_json_attr("device_group_attr_key", "device_group") g._set_json_attr("device_assign_map", {"stage1": 0, "stage2" : 1}, "dict_str_int") g._set_json_attr("device_copy_op", "cross_device_copy") g = g.apply("PlaceDevice") 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('device')[jnode_row_ptr[nindex["add2"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["add3"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["cast1"]]] == 0
def test_plan_memory(): x = sym.Variable('x', shape=(4, 2)) x2 = sym.add(x, x, name='addk') y = sym.reshape(x2, target=(2, 4), name="reshapek") y = sym.add(y, x2, name="add2") y = sym.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_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 test_place_device(): x = sym.Variable('x', device_group="stage1") y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") z = sym.add(y, y, device_group="stage2", name="add2") z = sym.add(z, sym.exp(y, device_group="stage2"), name="add3") g = graph.create(z) g._set_json_attr("device_group_attr_key", "device_group") g._set_json_attr("device_assign_map", {"stage1": 0, "stage2" : 1}, "dict_str_int") g._set_json_attr("device_copy_op", "cross_device_copy") g = g.apply("PlaceDevice") 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('device')[jnode_row_ptr[nindex["add2"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["add3"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["cast1"]]] == 0
def test_alter_func_return_none(): data = sym.Variable("data", shape=(1, 32, 512, 512)) pool1 = sym.global_max_pool2d(data, name="pool1", layout="NCHW") pool2 = sym.global_max_pool2d(pool1, name="pool2", layout="NCHW") relu = sym.relu(pool2, name="relu") g = graph.create(relu) g = g.apply("CorrectLayout") g = graph_attr.set_dtype_inputs(g, "float32") g = g.apply(["InferShape", "InferType"]) assert g.json_attr("layout") == ['NCHW', 'NCHW', 'NCHW', 'NCHW'] @reg.register_alter_op_layout("global_max_pool2d", level=100) def alter_global_max_pool2d_layout(attrs, inputs, tinfos): return None g = g.apply("AlterOpLayout") # alter func return none, nothing get replaced, # the layouts should remain the same assert g.json_attr("layout") == ['NCHW', 'NCHW', 'NCHW', 'NCHW']
def infer_variable_shapes(net, feed_dict): """Inference shape of all variables in the net. Parameters ---------- net : tf.Symbol The symbolic network containing all the variables. feed_dict : dict dict of placeholder to known shape Returns ------- Generator of (var, vname, vshape) Enables enumeration of variables in the net with corresponding name and shape. """ g = graph.create(net) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnode_row_ptr = jgraph["node_row_ptr"] jnodes = jgraph["nodes"] shape = [[]] * jnode_row_ptr[-1] nindex = {n['name']: i for i, n in enumerate(jnodes)} for k, v in feed_dict.items(): node_name = k.attr("name") shape[jnode_row_ptr[nindex[node_name]]] = v g._set_json_attr("shape", shape, "list_shape") g = g.apply("InferShape") shape = g.json_attr("shape") ret = {} for v in net.list_input_variables(): vname = v.attr("name") vshape = shape[jnode_row_ptr[nindex[vname]]] if len(vshape) == 0: raise ValueError("not sufficient information in feed_dict") yield (v, vname, vshape)
def infer_variable_shapes(net, feed_dict): """Inference shape of all variables in the net. Parameters ---------- net : tf.Symbol The symbolic network containing all the variables. feed_dict : dict dict of placeholder to known shape Returns ------- Generator of (var, vname, vshape) Enables enumeration of variables in the net with corresponding name and shape. """ g = graph.create(net) jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnode_row_ptr = jgraph["node_row_ptr"] jnodes = jgraph["nodes"] shape = [[]] * jnode_row_ptr[-1] nindex = {n['name']: i for i, n in enumerate(jnodes)} for k, v in feed_dict.items(): node_name = k.attr("name") shape[jnode_row_ptr[nindex[node_name]]] = v g._set_json_attr("shape", shape, "list_shape") g = g.apply("InferShape") shape = g.json_attr("shape") ret = {} for v in net.list_input_variables(): vname = v.attr("name") vshape = shape[jnode_row_ptr[nindex[vname]]] if len(vshape) == 0: raise ValueError("not sufficient information in feed_dict") yield (v, vname, vshape)
conv_kernel = sym.Variable("conv_kernel") conv_bias = sym.Variable("conv_bias") sparse_kernel = sym.Variable("sparse_kernel") #y = sym.conv2d_sparse(data=x, weight=conv_kernel, sparsity=sparse_kernel, bias=conv_bias, channels=1, kernel_size=(3,3), padding=(0,0), use_bias=True, out_layout='NCHW') y = sym.conv2d_sparse(data=x, weight=conv_kernel, sparsity=sparse_kernel, channels=1, kernel_size=(3, 3), padding=(0, 0), use_bias=False, out_layout='NCHW') # Test Graph compilation # Once the API is well-defined, this part will be OK g = graph.create(y) print("-------------Starts----------------") print(g.json()) print("-----------------------------------") print(g.ir()) print("--------------Ends-----------------") # Check computation def forward(x, conv_kernel, sparse_kernel, kernel, pad, stride, **args): ishape = x.shape return def backward(head_grads, x, conv_kernel, sparse_kernel, kernel, pad, stride, **args):
def grad(ys, xs, ys_grads): g = graph.create(ys) g._set_symbol_list_attr('grad_ys', ys) g._set_symbol_list_attr('grad_xs', xs) g._set_symbol_list_attr('grad_ys_out_grad', ys_grads) return g.apply('Gradient')
def test_graph_json_attr(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2,2)) g = graph.create(y) g._set_json_attr('ilist', [1,2,3], 'list_int') assert g.json_attr('ilist') == [1,2,3]
def test_graph_json_attr(): x = sym.Variable('x') y = sym.dense(data=x, name='fc', units=30) g = graph.create(y) g._set_json_attr('ilist', [1, 2, 3], 'list_int') assert g.json_attr('ilist') == [1, 2, 3]
def test_graph_json_attr(): x = sym.Variable('x') y = sym.conv2d(data=x, name='conv', stride=(2, 2)) g = graph.create(y) g._set_json_attr('ilist', [1, 2, 3], 'list_int') assert g.json_attr('ilist') == [1, 2, 3]
def test_graph_json_attr(): x = sym.Variable('x') y = sym.dense(data=x, name='fc', units=30) g = graph.create(y) g._set_json_attr('ilist', [1,2,3], 'list_int') assert g.json_attr('ilist') == [1,2,3]