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)