def verify_resize(shape, scale, method, layout):
if layout == "NHWC":
size = (shape[1] * scale, shape[2] * scale)
else:
size = (shape[2] * scale, shape[3] * scale)
x = relay.var("x", relay.TensorType(shape, "float32"))
size_var = relay.const(np.array(size).astype("float32"))
coord_trans = "asymmetric" if method == "nearest_neighbor" else "align_corners"
z = relay.image.resize2d(x,
size_var,
layout,
method,
coordinate_transformation_mode=coord_trans)
func = run_infer_type(relay.Function([x], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("image.resize2d")
x_data = np.random.uniform(low=-1, high=1,
size=shape).astype("float32")
ref_res = tvm.topi.testing.resize2d_python(x_data, (scale, scale),
layout, method, coord_trans)
def verify_upsampling3d(data_shape, scale_d_val, scale_h_val, scale_w_val,
dtype):
x = relay.var("x", relay.TensorType(data_shape, dtype))
scale_d = relay.const(scale_d_val)
scale_h = relay.const(scale_h_val)
scale_w = relay.const(scale_w_val)
z = relay.nn.upsampling3d(x, scale_d, scale_h, scale_w)
func = run_infer_type(relay.Function([x], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("nn.upsampling3d")
x_data = np.random.uniform(size=data_shape).astype(dtype)
ref_res = tvm.topi.testing.resize3d_python(
x_data,
(scale_d_val, scale_h_val, scale_w_val),
"NCDHW",
"nearest_neighbor",
"asymmetric",
)
verify_func(func2, [x_data], ref_res)
def verify_upsampling(data_shape, scale_h_val, scale_w_val, dtype):
x = relay.var("x", relay.TensorType(data_shape, dtype))
scale_h = relay.var("scale_h", relay.TensorType((), "float32"))
scale_w = relay.var("scale_w", relay.TensorType((), "float32"))
z = relay.nn.upsampling(x, scale_h, scale_w)
params = {
"scale_h": scale_h_val,
"scale_w": scale_w_val,
}
func = run_infer_type(relay.Function([x, scale_h, scale_w], z))
func2 = run_opt_pass(
run_opt_pass(func, transform.DynamicToStatic(), params),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("nn.upsampling")
x_data = np.random.uniform(size=data_shape).astype(dtype)
ref_res = tvm.topi.testing.resize2d_python(x_data,
(scale_h_val, scale_w_val),
"NCHW", "nearest_neighbor",
"asymmetric")
verify_func(func2, [x_data], ref_res)
def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
if strides:
if len(strides) == 1:
strides = strides * ndim
else:
strides = [1] * ndim
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(x_data, begin, end, strides, slice_mode)
data = [x_data, np.array(begin), np.array(end)]
begin = relay.const(begin, dtype=dtype)
end = relay.const(end, dtype=dtype)
if strides:
data.append(np.array(strides))
strides = relay.const(strides, dtype=dtype)
z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode)
else:
z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode)
func = relay.Function([x], z)
func = run_infer_type(func)
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
assert isinstance(func2.body, relay.Call)
assert func2.body.op == relay.op.get("strided_slice")
verify_func(func2, [x_data], ref_res)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices",
relay.TensorType(indices_shape, "int32"))
depth_var = relay.const(depth)
on_value_var = relay.var("on_value", relay.TensorType((), "int32"))
off_value_var = relay.var("off_value", relay.TensorType((), "int32"))
out = relay.one_hot(indices, on_value_var, off_value_var, depth_var,
axis, dtype)
params = {
"on_value": on_value,
"off_value": off_value,
}
func = relay.Function([indices, on_value_var, off_value_var], out)
func2 = run_opt_pass(
run_opt_pass(func, transform.DynamicToStatic(), params),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("one_hot")
indices_np = np.random.randint(0, depth,
size=indices_shape).astype("int32")
out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value,
depth, axis, dtype)
verify_func(func2, [indices_np], out_np)
def verify_resize(shape, scale, method, layout):
if layout == "NHWC":
size = (shape[1] * scale, shape[2] * scale)
else:
size = (shape[2] * scale, shape[3] * scale)
x = relay.var("x", relay.TensorType(shape, "float32"))
size_var = relay.const(np.array(size).astype("float32"))
z = relay.image.resize(x, size_var, layout, method, "align_corners")
func = run_infer_type(relay.Function([x], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("image.resize")
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
if method == "bilinear":
ref_res = tvm.topi.testing.bilinear_resize_python(x_data, size, layout)
else:
ref_res = tvm.topi.testing.upsampling_python(x_data, (scale, scale), layout)
verify_func(func2, [x_data], ref_res, rtol=1e-4, atol=1e-6)
def verify(
dshape,
begin_val,
end_val,
strides_val,
output,
slice_mode="end",
test_ref=True,
dtype="int32",
):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin_val = begin_val if begin_val else [0] * ndim
end_val = end_val if end_val else list(dshape)
if strides_val:
if len(strides_val) == 1:
strides_val = strides_val * ndim
else:
strides_val = [1] * ndim
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(
x_data, begin_val, end_val, strides_val, slice_mode)
data = [x_data, np.array(begin_val), np.array(end_val)]
begin = relay.var("begin", relay.TensorType((len(begin_val), ), dtype))
end = relay.var("end", relay.TensorType((len(end_val), ), dtype))
func_params = [x, begin, end]
if strides_val:
data.append(np.array(strides_val))
strides = relay.var("strides",
relay.TensorType((len(strides_val), ), dtype))
z = relay.strided_slice(x,
begin=begin,
end=end,
strides=strides,
slice_mode=slice_mode)
func_params.append(strides)
else:
z = relay.strided_slice(x,
begin=begin,
end=end,
slice_mode=slice_mode)
func = relay.Function(func_params, z)
params = {"begin": begin_val, "end": end_val, "strides": strides_val}
func = run_infer_type(func)
func2 = run_opt_pass(
run_opt_pass(func, transform.DynamicToStatic(), params),
transform.InferType())
assert isinstance(func2.body, relay.Call)
assert func2.body.op == relay.op.get("strided_slice")
verify_func(func2, [x_data], ref_res)
def verify_pad(data_shape, pad_width, pad_val, dtype):
x = relay.var("x", relay.TensorType(data_shape, dtype))
z = relay.nn.pad(x, relay.const(np.array(pad_width)), pad_val)
func = run_infer_type(relay.Function([x], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("nn.pad")
x_data = np.random.uniform(size=data_shape).astype(dtype)
ref_res = np.pad(x_data, pad_width, 'constant', constant_values=(((pad_val,)*2),) * len(data_shape))
verify_func(func2, [x_data], ref_res)
def test_dynamic_to_static_dynamic_if():
x = relay.var("x", relay.TensorType((2, 2), "int64"))
cond = relay.const(1)
iff = relay.If(cond, relay.reshape(x, [1, 4]), relay.reshape(x, (4, 1)))
func = relay.Function([x], iff)
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("reshape")
x_data = np.random.uniform(low=-1, high=1, size=(2, 2)).astype("int64")
verify_func(func2, [x_data], x_data.reshape(1, 4))
def verify_topk(k, axis, ret_type, is_ascend, dtype):
shape = (20, 100)
x = relay.var("x", relay.TensorType(shape, "float32"))
k_var = relay.const(k)
out = relay.topk(x, k_var, axis, ret_type, is_ascend, dtype)
if isinstance(out, relay.expr.TupleWrapper):
out = out.astuple()
func = relay.Function([x], out)
np_data = np.random.uniform(size=shape).astype("float32")
if is_ascend:
np_indices = np.argsort(np_data, axis=axis)
else:
np_indices = np.argsort(-np_data, axis=axis)
kk = k if k >= 1 else shape[axis]
if axis == 0:
np_indices = np_indices[:kk, :]
np_values = np.zeros(np_indices.shape).astype("float32")
for i in range(shape[1]):
np_values[:, i] = np_data[np_indices[:, i], i]
else:
np_indices = np_indices[:, :kk]
np_values = np.zeros(np_indices.shape).astype("float32")
for i in range(shape[0]):
np_values[i, :] = np_data[i, np_indices[i, :]]
np_indices = np_indices.astype(dtype)
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("topk")
for target, ctx in tvm.testing.enabled_targets():
if "llvm" not in target:
continue
for kind in ["graph", "vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func2)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
op_res = intrp.evaluate()(np_data)
if ret_type == "both":
tvm.testing.assert_allclose(op_res[0].asnumpy(), np_values)
tvm.testing.assert_allclose(op_res[1].asnumpy(),
np_indices)
elif ret_type == "values":
tvm.testing.assert_allclose(op_res.asnumpy(), np_values)
else:
tvm.testing.assert_allclose(op_res.asnumpy(), np_indices)
def verify_full(fill_value, fill_shape, dtype):
x = relay.var("x", relay.scalar_type(dtype))
y = relay.var("y", relay.TensorType(fill_shape, 'int64'))
z = relay.full(x, relay.shape_of(y), dtype)
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("full")
ref_res = np.full(fill_shape, fill_value).astype(dtype)
y_data = np.random.uniform(low=-1, high=1, size=fill_shape).astype('int64')
verify_func(func2, [fill_value, y_data], ref_res)
def verify_ones_zeros(shape, dtype):
for op, ref in [(relay.zeros, np.zeros), (relay.ones, np.ones)]:
x = relay.var("x", relay.TensorType(shape, dtype))
y = op(relay.shape_of(x), dtype)
func = run_infer_type(relay.Function([x], y))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Constant)
assert zz.checked_type == relay.ty.TensorType(shape, dtype)
x_data = np.random.uniform(low=1, high=1, size=shape)
ref_res = ref(x_data.shape)
verify_func(func2, [x_data], ref_res)
def verify_tile(shape, reps, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType(reps, "float32"))
z = relay.tile(x, relay.shape_of(y))
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("tile")
assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
y_data = np.random.uniform(low=-1, high=1, size=reps).astype("float32")
ref_res = np.tile(x_data, reps)
verify_func(func2, [x_data, y_data], ref_res)
def verify_sparse_to_dense(sparse_indices, sparse_values, default_value,
output_shape, xpected):
sparse_indices_data = np.array(sparse_indices)
sparse_values_data = np.array(sparse_values)
default_value_data = np.array(default_value)
output_shape_data = np.array(output_shape)
a = relay.var(
"a",
relay.TensorType(sparse_indices_data.shape,
str(sparse_indices_data.dtype)))
b = relay.var(
"b",
relay.TensorType(sparse_values_data.shape,
str(sparse_values_data.dtype)))
output_shape_const = relay.const(output_shape_data)
if default_value is None:
args = [a, b]
d = relay.sparse_to_dense(a, output_shape_const, b)
else:
c = relay.var(
"c",
relay.TensorType(default_value_data.shape,
str(default_value_data.dtype)))
args = [a, b, c]
d = relay.sparse_to_dense(a, output_shape_const, b, c)
zz = run_infer_type(d)
assert len(zz.checked_type.shape) == len(output_shape)
func = relay.Function(args, d)
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
assert isinstance(func2.body, relay.Call)
assert func2.body.op == relay.op.get("sparse_to_dense")
if default_value is None:
arguments = [sparse_indices_data, sparse_values_data]
else:
arguments = [
sparse_indices_data, sparse_values_data, default_value_data
]
verify_func(func2, arguments, xpected)
def verify_reshape(shape, newshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType(newshape, "float32"))
z1 = relay.reshape(x, relay.shape_of(y))
z2 = relay.reshape(z1, relay.shape_of(x))
z3 = relay.reshape(z2, relay.shape_of(z1))
z4 = relay.reshape(z3, relay.shape_of(z2))
func = run_infer_type(relay.Function([x, y], z4))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("reshape")
assert "newshape=" in zz.astext()
assert zz.checked_type == relay.ty.TensorType(shape, "float32")
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
y_data = np.random.uniform(low=-1, high=1, size=newshape).astype("float32")
verify_func(func2, [x_data, y_data], x_data)
def partition_for_bnns(mod, params=None):
"""Partition the graph greedily offloading supported
operators to BNNS.
Parameters
----------
mod : Module
The module to run passes on.
params : Optional[Dict[str, NDArray]]
Constant input parameters.
Returns
-------
ret : annotated and partitioned module.
"""
if params:
mod["main"] = bind_params_by_name(mod["main"], params)
seq = tvm.transform.Sequential(
[
transform.InferType(),
transform.FoldConstant(),
transform.FoldScaleAxis(),
transform.DynamicToStatic(),
transform.AlterOpLayout(),
# TODO(apeskov): WA. AlterOpLayout call lead to constants shape transformation
# Some expand_dims op may appears after constants. It breaks BNNS fusing.
# So we have to call FoldConstant right before bnns composite passes.
transform.FoldConstant(),
transform.MergeComposite(get_pattern_table("bnns")),
transform.AnnotateTarget("bnns"),
# If you no need in per layer performance statistic you can
# uncomment next line
# transform.MergeCompilerRegions(),
transform.PartitionGraph(),
]
)
return seq(mod)
