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]
