def test_any_concat(): x = relay.var('x', shape=(relay.Any(), 2), dtype="float32") y = relay.var('y', shape=(1, 2), dtype="float32") xx = x - relay.expr.const(3.0) yy = y * relay.expr.const(5.0) z = relay.op.concatenate([xx, yy], axis=0) mod = tvm.IRModule() mod["main"] = relay.Function([x, y], z) x_np = np.random.uniform(size=(3, 2)).astype('float32') y_np = np.random.uniform(size=(1, 2)).astype('float32') ref = np.concatenate([x_np - 3.0, y_np * 5.0], axis=0) check_result([x_np, y_np], mod, ref)
def test_any_crop_and_resize(): verify_any_crop_and_resize( data_shape=(1, 234, 234, 256), boxes_shape=(relay.Any(), 4), box_indices_shape=(relay.Any(),), crop_size=(14, 14), layout="NHWC", static_boxes=(128, 4), static_box_indices_shape=(128,), ref_out_shape=(128, 14, 14, 256), ) verify_any_crop_and_resize( data_shape=(1, 256, 234, 234), boxes_shape=(relay.Any(), 4), box_indices_shape=(relay.Any(),), crop_size=(14, 14), layout="NCHW", static_boxes=(128, 4), static_box_indices_shape=(128,), ref_out_shape=(128, 256, 14, 14), )
def test_robust_imputer(): st_helper = SklearnTestHelper() data = np.array( [[4, 5, np.nan, 7], [0, np.nan, 2, 3], [8, 9, 10, 11], [np.inf, 13, 14, 15]], dtype=np.float32, ) ri = RobustImputer(dtype=None, strategy="constant", fill_values=np.nan, mask_function=None) ri.fit(data) dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, ri, dshape, data)
def test_simple_imputer(): st_helper = SklearnTestHelper() data = np.array( [[4, 5, np.nan, 7], [0, np.nan, 2, 3], [8, 9, 10, 11], [np.nan, 13, 14, 15]], dtype=np.float32, ) imp_mean = SimpleImputer(missing_values=np.nan, strategy="median") imp_mean.fit(data) dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, imp_mean, dshape, data)
def test_robust_missing_indicator(): st_helper = SklearnTestHelper() data = np.array( [[4, 5, np.nan, 7], [0, np.nan, 2, 3], [8, 9, 10, 11], [np.inf, 13, 14, 15]], dtype=np.float32, ) rmi = RobustMissingIndicator() rmi.fit(data) dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, rmi, dshape, data)
def test_unique(target, dev): dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() [unique, _, _, num_unique] = relay.unique(x, is_sorted=True) mod["main"] = relay.Function([x], relay.op.strided_slice(unique, begin=[0], end=num_unique)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.unique(x_np) check_mod(target, dev, mod, x_np, res_np)
def test_dynamic_channels(): # Compile simulated quantize once but support either per-channel or scalar params. data = np.random.uniform(low=-64, high=64, size=[2, 5]).astype("float32") # Test scalar qnn params. scale_np = np.asarray([0.5]).astype("float32") zp_np = np.asarray([127]).astype("int32") dtype_np = np.int32(SQNN_DTYPE_TO_CODE["uint8"]) quant_args = {"out_zero_point": zp_np[0], "out_scale": scale_np[0]} q_out = quantize_test_driver( in_dtype="float32", quant_args=quant_args, axis=0, out_dtype="uint8", in_data=data, ) # Create variables with undefined shape and run with scalar inputs. input_data = relay.var("input_data", shape=data.shape, dtype="float32") scale = relay.var("scale", shape=[relay.Any()], dtype="float32") zp = relay.var("zp", shape=[relay.Any()], dtype="int32") dtype = relay.var("dtype", shape=[]) vm = build_simulated_quantize(input_data, scale, zp, dtype, axis=0) sim_q_out = vm.invoke("main", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np) allclose_with_rounding(sim_q_out.asnumpy(), q_out) # Now get the perchannel quantize output and compare without recompiling. scale_np = np.array([0.5, 0.25]).astype("float32") zp_np = np.array([127, 123]).astype("int32") # Get the reference quantize output. quant_args = {"out_zero_point": zp_np, "out_scale": scale_np} q_out = quantize_test_driver( in_dtype="float32", quant_args=quant_args, axis=0, out_dtype="uint8", in_data=data, ) # Run the simulated quantize without recompiling and confirm results match. sim_q_out = vm.invoke("main", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np) allclose_with_rounding(sim_q_out.asnumpy(), q_out)
def test_any_conv2d_NCHWc(): verify_any_conv2d_NCHWc( (relay.Any(), 8, relay.Any(), relay.Any(), 8), (8, 8, 3, 3, 8, 8), (1, 1), (1, 1), (1, 1), "NCHW8c", "OIHW8i8o", "NCHW8c", (1, 8, 224, 224, 8), (1, 8, 224, 224, 8), ) verify_any_conv2d_NCHWc( (relay.Any(), 8, relay.Any(), relay.Any(), 8), (8, 8, 3, 3, 8, 8), (1, 1), (1, 1), (2, 2), "NCHW8c", "OIHW8i8o", "NCHW8c", (1, 8, 224, 224, 8), (1, 8, 222, 222, 8), )
def test_dynamic_bcast(): import pytest import numpy as np import tvm from tvm.runtime import vm as _vm from tvm.relay import vm as rly_vm from tvm import relay from tvm.relay.scope_builder import ScopeBuilder from tvm.relay import transform from tvm.relay.prelude import Prelude from tvm.relay import testing def create_exec(f, target="llvm", params=None): if isinstance(f, relay.Expr): mod = tvm.IRModule() mod["main"] = f executable = rly_vm.compile(mod, target=target, params=params) return executable else: assert isinstance(f, tvm.IRModule), "expected mod as tvm.IRModule" executable = rly_vm.compile(f, target=target, params=params) return executable def get_serialized_output(mod, *data, params=None, target="llvm", ctx=tvm.cpu()): exe = create_exec(mod, target, params=params) code, lib = exe.save() des_exec = _vm.Executable.load_exec(code, lib) des_vm = _vm.VirtualMachine(des_exec, ctx) result = des_vm.run(*data) print(result) return result dtype = "float32" x = relay.var("x", shape=(1, 2), dtype=dtype) y = relay.var("y", shape=(relay.Any(), 2), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x, y], relay.add(x, y)) x_data = np.random.uniform(size=(1, 2)).astype(dtype) y_data = np.random.uniform(size=(4, 2)).astype(dtype) res_np = np.add(x_data, y_data) for target, ctx in testing.enabled_targets(): res = get_serialized_output(mod, *(x_data, y_data), target=target, ctx=ctx) tvm.testing.assert_allclose(res.asnumpy(), res_np)
def test_dyn_upsampling3d_infer_type_const(): n, c, d, h, w = ( te.size_var("n"), te.size_var("c"), te.size_var("d"), te.size_var("h"), te.size_var("w"), ) data = relay.var("data", relay.TensorType((n, c, d, h, w), "int8")) scale_d = relay.Var("scale_h", relay.TensorType((), "float32")) scale_w = relay.Var("scale_w", relay.TensorType((), "float32")) z = relay.nn.upsampling3d(data, scale_d, 2.0, scale_w, layout="NCDHW", method="trilinear") zz = run_infer_type(z) assert zz.checked_type == relay.TensorType( (n, c, relay.Any(), relay.Any(), relay.Any()), "int8")
def test_pca(): st_helper = SklearnTestHelper() pca = PCA(n_components=2) rpca = RobustPCA() data = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]], dtype=np.float32) pca.fit(data) rpca.robust_pca_ = pca dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, rpca, dshape, data) tSVD = TruncatedSVD(n_components=5, n_iter=7, random_state=42) data = sparse_random(100, 100, density=0.01, format="csr", dtype="float32", random_state=42).toarray() tSVD.fit(data) rpca.robust_pca_ = tSVD dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, rpca, dshape, data)
def test_any_concat(): x = relay.var('x', shape=(relay.Any(), 2), dtype="float32") y = relay.var('y', shape=(1, 2), dtype="float32") z = relay.op.concatenate([x, y], axis=0) mod = relay.module.Module() mod["main"] = relay.Function([x, y], z) x_np = np.random.uniform(size=(3, 2)).astype('float32') y_np = np.random.uniform(size=(1, 2)).astype('float32') ref = np.concatenate([x_np, y_np], axis=0) for kind in ["debug", "vm"]: ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm") result = ex.evaluate()(x_np, y_np) tvm.testing.assert_allclose(result.asnumpy(), ref)
def _test_quantile_transformer(shape, n_quantiles): from sklearn.preprocessing import QuantileTransformer st_helper = SklearnTestHelper() rng = np.random.RandomState(0) data = np.sort(rng.normal(loc=0.5, scale=0.25, size=shape), axis=0) qt = QuantileTransformer(n_quantiles=n_quantiles, random_state=0) qt.fit_transform(data) dshape = (relay.Any(), len(data[0])) _test_model_impl(st_helper, qt, dshape, data.astype("float32"))
def test_any_split(): verify_any_split((relay.Any(), 4), 2, 1, (9, 4), [(9, 2), (9, 2)]) verify_any_split((relay.Any(), relay.Any()), 2, 1, (9, 4), [(9, 2), (9, 2)]) verify_any_split((relay.Any(), 12), (1, 4, 8), 1, (7, 12), [(7, 1), (7, 3), (7, 4)]) verify_any_split((relay.Any(), relay.Any()), (1, 4, 8), 1, (7, 12), [(7, 1), (7, 3), (7, 4)])
def test_mixed_input_type(): mod = tvm.IRModule() dtype = "float32" static_data_shape = (9, 4) data_shape = (relay.Any(), 4) tensor_type = relay.TensorType(data_shape, dtype) tuple_type = relay.TupleType([tensor_type, tensor_type]) data0 = relay.var("d0", type_annotation=relay.TupleType([tuple_type, tensor_type])) data1 = relay.var("d1", shape=(relay.Any(), 4), dtype=dtype) data_tuple = relay.expr.TupleWrapper(data0, 2) nested_data_tuple = relay.expr.TupleWrapper(data_tuple[0], 2) y = nested_data_tuple[1] * data_tuple[1] + data1 mod["main"] = relay.Function([data0, data1], y) data_np0 = np.random.uniform(size=static_data_shape).astype(dtype) data_np1 = np.random.uniform(size=static_data_shape).astype(dtype) ref_out_shape = (9, 4) check_result( [[[data_np0, data_np0], data_np0], data_np1], mod, ref_out_shape, assert_shape=True, only_vm=True, )
def verify_pad_default_fill(dshape, pad_width, dtype): x = relay.var("x", relay.TensorType(dshape, dtype)) ndim = len(dshape) pad_width_var = relay.var("pad_width_var", relay.TensorType((ndim, 2), 'int64')) y = relay.nn.pad(x, pad_width_var) yy = run_infer_type(y) assert yy.checked_type == relay.ty.TensorType((relay.Any(),) * ndim, dtype) func = relay.Function([x, pad_width_var], y) data = np.random.uniform(size=dshape).astype(dtype) ref_res = np.pad(data, pad_width) pad_width = np.array(pad_width).astype('int64') verify_func(func, [data, pad_width], ref_res)
def test_pushconstants(): if not tvm.testing.device_enabled("vulkan"): return # Three 32 bit pushconstants: any_dim, stride, stride dtype = "float32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.sqrt(x)) x_np = np.random.uniform(size=(10, )).astype(dtype) res_np = np.sqrt(x_np) check_mod(mod, x_np, res_np) # One 64 bit and one 32 bit constants dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x], relay.argsort(x)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.argsort(x_np) check_mod(mod, x_np, res_np)
def test_mixed_input_type(): mod = tvm.IRModule() dtype = "float32" static_data_shape = (9, 4) data_shape = (relay.Any(), 4) tensor_type = relay.TensorType(data_shape, dtype) tuple_type = relay.TupleType([tensor_type, tensor_type]) data0 = relay.var("d0", type_annotation=relay.TupleType( [tuple_type, tensor_type])) data1 = relay.var("d1", shape=(relay.Any(), 4), dtype=dtype) data_tuple = relay.expr.TupleWrapper(data0, 2) nested_data_tuple = relay.expr.TupleWrapper(data_tuple[0], 2) y = nested_data_tuple[1] * data_tuple[1] + data1 mod["main"] = relay.Function([data0, data1], y) data_np0 = np.random.uniform(size=static_data_shape).astype(dtype) data_np1 = np.random.uniform(size=static_data_shape).astype(dtype) ref_out_shape = (9, 4) for kind in ["vm"]: ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm") result = ex.evaluate()([[data_np0, data_np0], data_np0], data_np1) assert result.asnumpy().shape == ref_out_shape, \ "Shape mismatch: expect %s but got %s." % (str(ref_out_shape), str(result.asnumpy().shape))
def test_take_grad(): data_dtype = relay.TensorType((3, 4, 5), "float64") data = relay.var("data", data_dtype) indices = relay.var("indices", relay.TensorType((relay.Any(),), "int32")) inputs = [_np_randn_from_type(data_dtype, scale=1e-5), np.array([1, 2], dtype="int32")] test_inputs = [inputs[0]] # take on axis fwd_func = relay.Function([data, indices], relay.take(data, indices, axis=1)) check_grad(fwd_func, inputs=inputs, test_inputs=test_inputs) # take on flattened fwd_func = relay.Function([data, indices], relay.take(data, indices, axis=None)) check_grad(fwd_func, inputs=inputs, test_inputs=test_inputs)
def test_shape_func_nested_function(): @tvm.register_func("relay.ext.test2") def relay_ext_test(func): return None data_shape = (relay.Any(), 16) weight_shape = (relay.Any(), 16) dense = relay.nn.dense(relay.var("data", shape=data_shape), relay.var("weight", shape=weight_shape)) mod = tvm.IRModule.from_expr(dense) patterns = [("test.dense", is_op("nn.dense")(wildcard(), wildcard()))] passes = tvm.transform.Sequential([ relay.transform.MergeComposite(patterns), relay.transform.AnnotateTarget(["test2"]), relay.transform.PartitionGraph(), ]) mod = passes(mod) compiler = VMCompiler() compiler.lower(mod, "llvm")
def verify_zeros_ones(shape, dtype): for op, ref in [(relay.zeros, np.zeros), (relay.ones, np.ones)]: rank = len(shape) dyn_shape = relay.Var("shape", relay.ty.TensorType((rank, ), "int64")) y = op(dyn_shape, dtype) yy = run_infer_type(y) assert yy.checked_type == relay.ty.TensorType( (relay.Any(), ) * rank, dtype) func = relay.Function([dyn_shape], y) ref_res = ref(shape, dtype) verify_func(executor_kind, func, [np.array(shape).astype("int64")], ref_res.astype("int64"))
def test_tuple_get_item(): mod = tvm.IRModule() dtype = "float32" static_data_shape = (9, 4) data_shape = (relay.Any(), 4) indices_or_sections = 2 axis = 1 data = relay.var("data", shape=data_shape, dtype=dtype) y = relay.split(data, indices_or_sections, axis) y = relay.expr.TupleGetItem(y.astuple(), 0) mod["main"] = relay.Function([data], y) data_np = np.random.uniform(size=static_data_shape).astype(dtype) ref_out_shape = (9, 2) check_result([data_np], mod, ref_out_shape, assert_shape=True)
def test_any_conv2d_transpose_nchw(): verify_any_conv2d_transpose_nchw( (relay.Any(), 64, 224, 224), (64, 192, 3, 3), (1, 1), (1, 1), (1, 1), 1, (2, 64, 224, 224), (2, 192, 224, 224), (0, 0), ) verify_any_conv2d_transpose_nchw( (relay.Any(), 32, 224, 224), (32, 64, 3, 3), (2, 2), (1, 1), (1, 1), 1, (1, 32, 224, 224), (1, 64, 448, 448), (1, 1), )
def test_fuse_dynamic_squeeze_slice_take(): input_data = [ np.random.random([1, 2, 4]).astype("float32"), np.array([0]).astype("int64"), ] x = relay.var("p0107", shape=(relay.Any(), relay.Any(), 4), dtype="float32") take_val = relay.var("p166", shape=(relay.Any(),), dtype="int64") squeeze = relay.op.squeeze(x, axis=[0]) strided_slice = relay.op.strided_slice( squeeze, begin=[0, 0], end=[15130, 9223372036854775807], strides=[1, 1] ) take = relay.op.take(strided_slice, take_val, axis=0) mod = tvm.IRModule.from_expr(take) result = relay.create_executor("vm", mod=mod, device=tvm.cpu(), target="llvm").evaluate()( *input_data ) np_result = np.squeeze(input_data[0][:, input_data[1][0], :], axis=0) assert np.allclose(result.numpy(), np_result)
def test_dense_dynamic(): data_shape = (relay.Any(), K) weight_shape = (relay.Any(), K) if has_cublas(): # TVM native fp16 dense (without tensorcore), using fp16 accum, seems to have accuracy issues # Use cublas as a reference verify_dense( get_dense_with_shape(data_shape, weight_shape), M, N, K, ref_target="cuda -libs=cublas", ) verify_dense( get_dense_with_shape(data_shape, weight_shape, out_dtype="float32"), M, N, K, atol=1e-4, rtol=1e-4, )
def test_dynamic_dtype(): # Compile simulated quantize once but support any type of quantization. data = np.random.uniform(low=-64, high=64, size=[2, 5]).astype("float32") # Test scalar float32 to uint8. scale_np = np.asarray([0.5]).astype("float32") zp_np = np.asarray([127]).astype("int32") dtype_np = np.int32(SQNN_DTYPE_TO_CODE["uint8"]) quant_args = {"out_zero_point": zp_np[0], "out_scale": scale_np[0]} q_out = quantize_test_driver( in_dtype="float32", quant_args=quant_args, axis=-1, out_dtype="uint8", in_data=data, ) # Create variables with undefined shape and run with scalar inputs. input_data = relay.var("input_data", shape=data.shape, dtype="float32") scale = relay.var("scale", shape=[relay.Any()], dtype="float32") zp = relay.var("zp", shape=[relay.Any()], dtype="int32") dtype = relay.var("dtype", shape=[]) vm = build_simulated_quantize(input_data, scale, zp, dtype) sim_q_out = vm.invoke("main", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np) np.testing.assert_equal(sim_q_out.asnumpy(), q_out) # Now test float32 to int32 compilation. # Get the reference quantize output. q_out = quantize_test_driver( in_dtype="float32", quant_args=quant_args, axis=-1, out_dtype="int32", in_data=data, ) # Run the simulated quantize without recompiling and confirm results match. dtype_np = np.int32(SQNN_DTYPE_TO_CODE["int32"]) sim_q_out = vm.invoke("main", input_data=data, scale=scale_np, zp=zp_np, dtype=dtype_np) np.testing.assert_equal(sim_q_out.asnumpy(), q_out)
def test_unique(): if not tvm.testing.device_enabled("vulkan"): return dtype = "int32" x = relay.var("x", shape=(relay.Any(), ), dtype=dtype) mod = tvm.IRModule() [unique, _, num_unique] = relay.unique(x, is_sorted=True) mod["main"] = relay.Function([x], relay.op.strided_slice(unique, begin=[0], end=num_unique)) x_np = np.random.randint(0, high=10, size=(10, )).astype(dtype) res_np = np.unique(x_np) check_mod(mod, x_np, res_np)
def test_dynamic_bcast(): dtype = "float32" x = relay.var("x", shape=(relay.Any(), 2), dtype=dtype) y = relay.var("y", shape=(3, 2), dtype=dtype) mod = tvm.IRModule() mod["main"] = relay.Function([x, y], relay.add(x, y)) x_data = np.random.uniform(size=(1, 2)).astype(dtype) y_data = np.random.uniform(size=(3, 2)).astype(dtype) res_np = np.add(x_data, y_data) for target, ctx in testing.enabled_targets(): res = get_serialized_output(mod, *(x_data, y_data), target=target, ctx=ctx) tvm.testing.assert_allclose(res.asnumpy(), res_np)
def verify_pad(dshape, pad_width, pad_val, dtype): x = relay.var("x", relay.TensorType(dshape, dtype)) ndim = len(dshape) pad_width_var = relay.var("pad_width_var", relay.TensorType((ndim, 2), "int64")) pad_val_var = relay.var("pad_val_var", relay.TensorType((), dtype)) y = relay.nn.pad(x, pad_width_var, pad_val_var) yy = run_infer_type(y) assert yy.checked_type == relay.ty.TensorType((relay.Any(),) * ndim, dtype) func = relay.Function([x, pad_width_var, pad_val_var], y) data = np.random.uniform(size=dshape).astype(dtype) ref_res = np.pad(data, pad_width, "constant", constant_values=(((pad_val,) * 2),) * ndim) pad_width = np.array(pad_width).astype("int64") verify_func(func, [data, pad_width, np.array(pad_val).astype(dtype)], ref_res)
def test_vm_shape_of(): if not tvm.testing.device_enabled("cuda") or not tvm.cuda(0).exist: return mod = tvm.IRModule() data_shape = (relay.Any(), ) x = relay.var("x", shape=data_shape) y = relay.op.vm.shape_of(x) mod["main"] = relay.Function([x], y) ca = context_analysis(mod, tvm.cuda()) main = mod["main"] cpu_dev = tvm.cpu().device_type gpu_dev = tvm.cuda().device_type assert main.params[0] in ca and ca[main.params[0]][0].value == gpu_dev assert main.body in ca and ca[main.body][0].value == cpu_dev