def test_concatenate_split():
x = sym.Variable('x')
y = sym.Variable('y')
y = sym.concatenate(x, y)
assert y.list_input_names() == ['x', 'y']
z = sym.split(y, indices_or_sections=10)
assert len(z.list_output_names()) == 10
z = sym.split(y, indices_or_sections=[10, 20])
assert len(z.list_output_names()) == 3
def test_concatenate_split():
x = sym.Variable('x')
y = sym.Variable('y')
y = sym.concatenate(x, y)
assert y.list_input_names() == ['x', 'y']
z = sym.split(y, indices_or_sections=10)
assert len(z.list_output_names()) == 10
z = sym.split(y, indices_or_sections=[10, 20])
assert len(z.list_output_names()) == 3
def test_split():
x1 = sym.Variable("x", shape=(10, 20))
z = sym.split(x1, indices_or_sections=[11], name="y")
sdict = infer_shape(z)
assert (sdict["y"][0] == [10, 11])
assert (sdict["y"][1] == [10, 9])
z = sym.split(x1, indices_or_sections=2, name="y")
sdict = infer_shape(z)
assert (sdict["y"][0] == [10, 10])
assert (sdict["y"][1] == [10, 10])
def test_split():
x1 = sym.Variable("x", shape=(10, 20))
z = sym.split(x1, indices_or_sections=[11], name="y")
sdict = infer_shape(z)
assert(sdict["y"][0] == [10, 11])
assert(sdict["y"][1] == [10, 9])
z = sym.split(x1, indices_or_sections=2, name="y")
sdict = infer_shape(z)
assert(sdict["y"][0] == [10, 10])
assert(sdict["y"][1] == [10, 10])
def verify_split(ishape, indices_or_sections, axis):
x = sym.Variable("x", shape=ishape)
y = sym.split(x, indices_or_sections=indices_or_sections, axis=axis)
def forward(x):
return np.split(x, indices_or_sections, axis=axis)
check_function(y, forward)
def verify_split(ishape, indices_or_sections, axis):
x = sym.Variable("x", shape=ishape)
y = sym.split(x, indices_or_sections=indices_or_sections, axis=axis)
def forward(x):
return np.split(x, indices_or_sections, axis=axis)
check_function(y, forward)
def test_split():
x = sym.Variable("x", shape=(10, 20))
y = sym.split(x, indices_or_sections=[11], name="y")
g, ldict = correct_layout(y, "HW")
assert (ldict["x"][0] == "HW")
assert (ldict["y"][0] == "__undef__")
# second pass will insert layout transform
_, ldict = correct_layout(g, "HW16w")
assert (ldict["x"][0] == "HW16w")
assert (ldict["x_HW"][0] == "HW")
assert (ldict["y"][0] == "__undef__")
def test_num_outputs():
x = sym.Variable('x')
z = sym.split(x, indices_or_sections=5, axis=1)
shape = (10, 10)
dtype = tvm.float32
nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
params = {"x": nx}
graph, lib, params = nnvm.compiler.build(
z, "llvm", shape={"x": nx.shape}, params=params)
m = graph_runtime.create(graph, lib, tvm.cpu(0))
assert m.get_num_outputs() == 5
def test_split():
x = sym.Variable("x", shape=(10, 20))
y = sym.split(x, indices_or_sections=[11], name="y")
g, ldict = correct_layout(y, "HW")
assert(ldict["x"][0] == "HW")
assert(ldict["y"][0] == "__undef__")
# second pass will insert layout transform
_, ldict = correct_layout(g, "HW16w")
assert(ldict["x"][0] == "HW16w")
assert(ldict["x_HW"][0] == "HW")
assert(ldict["y"][0] == "__undef__")
def test_num_outputs():
x = sym.Variable('x')
z = sym.split(x, indices_or_sections=5, axis=1)
shape = (10, 10)
dtype = tvm.float32
nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
params = {"x": nx}
graph, lib, params = nnvm.compiler.build(
z, "llvm", shape={"x": nx.shape}, params=params)
m = graph_runtime.create(graph, lib, tvm.cpu(0))
assert m.get_num_outputs() == 5
def verify_split(ishape, indices_or_sections, axis):
x = sym.Variable("x")
y = sym.split(x, indices_or_sections=indices_or_sections, axis=axis)
dtype = "float32"
x_np = np.random.uniform(size=ishape).astype(dtype)
res = np.split(x_np, indices_or_sections, axis=axis)
for target, ctx in ctx_list():
# set input
graph, lib, _ = nnvm.compiler.build(y, target, {"x": ishape})
m = graph_runtime.create(graph, lib, ctx)
m.run(x=x_np)
for i, arr in enumerate(res):
out = m.get_output(i, tvm.nd.empty(arr.shape))
np.testing.assert_allclose(out.asnumpy(), arr, atol=1e-5, rtol=1e-5)
def verify_split(ishape, indices_or_sections, axis):
x = sym.Variable("x")
y = sym.split(x, indices_or_sections=indices_or_sections, axis=axis)
dtype = "float32"
x_np = np.random.uniform(size=ishape).astype(dtype)
res = np.split(x_np, indices_or_sections, axis=axis)
for target, ctx in ctx_list():
# set input
graph, lib, _ = nnvm.compiler.build(y, target, {"x": ishape})
m = graph_runtime.create(graph, lib, ctx)
m.run(x=x_np)
for i, arr in enumerate(res):
out = m.get_output(i, tvm.nd.empty(arr.shape))
np.testing.assert_allclose(out.asnumpy(), arr, atol=1e-5, rtol=1e-5)
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")
# split here
convs = sym.split(conv, indices_or_sections=2)
relus = [sym.relu(x, name="relu") for x in convs]
relu = sym.concatenate(*relus)
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", level=100)
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"]
