def test_no_zero_point(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (2, 1, 2, 4) data_dtype = "uint8" kernel_shape = (3, 1, 2, 2) kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=0, kernel_zero_point=0, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype) # int8 input data_shape = (2, 1, 2, 4) data_dtype = "int8" kernel_shape = (3, 1, 2, 2) kernel_dtype = "int8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=0, kernel_zero_point=0, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype)
def test_alter_layout_nhwc_nchw_arm(): """ Check NHWC to NHCW conversion for a small sequence of ops.""" def alter_conv2d(attrs, inputs, tinfos, out_type): import topi with tvm.target.create("llvm -device=arm_cpu"): return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, out_type) # Check NHWC conversion. def before_nhwc(): x = relay.var("x", shape=(1, 56, 56, 64)) weight1 = relay.var('weight1', shape=(3, 3, 64, 64)) weight2 = relay.var('weight2', shape=(3, 3, 64, 64)) y = relay.nn.conv2d(x, weight1, channels=64, kernel_size=(3, 3), data_layout='NHWC', kernel_layout='HWIO') y = relay.nn.relu(y) y = relay.nn.avg_pool2d(y, pool_size=(1, 1), layout='NHWC') y = relay.nn.conv2d(y, weight2, channels=64, kernel_size=(3, 3), data_layout='NHWC', kernel_layout='HWIO') y = relay.nn.relu(y) y = relay.Function(analysis.free_vars(y), y) return y def expected_nhwc(): x = relay.var("x", shape=(1, 56, 56, 64)) weight1 = relay.var('weight1', shape=(3, 3, 64, 64)) weight2 = relay.var('weight2', shape=(3, 3, 64, 64)) y = relay.layout_transform(x, "NHWC", "NCHW") weight1 = relay.layout_transform(weight1, "HWIO", "OIHW") weight2 = relay.layout_transform(weight2, "HWIO", "OIHW") y = relay.nn.conv2d(y, weight1, channels=64, kernel_size=(3, 3)) y = relay.nn.relu(y) y = relay.nn.avg_pool2d(y, pool_size=(1, 1)) y = relay.nn.conv2d(y, weight2, channels=64, kernel_size=(3, 3)) y = relay.nn.relu(y) y = relay.layout_transform(y, "NCHW", "NHWC") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before_nhwc() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected_nhwc(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_alter_layout_dual_path(): """ Test alternating the layout with two outputs. One path continues to use the new layout while one path fall backs to old layout. """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1") weight2 = relay.var("weight2") y = relay.nn.conv2d(x, weight1, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) y1 = relay.nn.conv2d(y, weight2, channels=32, kernel_size=(3, 3), padding=(1, 1)) y1 = relay.nn.relu(y1) y2 = relay.nn.batch_flatten(y) ret = relay.Tuple([y1, y2]) y = relay.Function(analysis.free_vars(ret), ret) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs["data_layout"] = "NCHW16c" return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1") weight2 = relay.var("weight2") y = relay.layout_transform(x, "NCHW", "NCHW16c") y = relay.nn.conv2d( y, weight1, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c" ) y = relay.nn.relu(y) y1 = relay.nn.conv2d( y, weight2, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c" ) y1 = relay.nn.relu(y1) y1 = relay.layout_transform(y1, "NCHW16c", "NCHW") y2 = relay.layout_transform(y, "NCHW16c", "NCHW") y2 = relay.nn.batch_flatten(y2) ret = relay.Tuple([y1, y2]) y = relay.Function(analysis.free_vars(ret), ret) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_layout(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (2, 2, 4, 4) # NHWC data_dtype = "uint8" kernel_shape = (2, 2, 4, 3) # HWIO kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=5, kernel_zero_point=3, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(1, 1), data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype) # NHWC and HWOI layout. Used in depthwise conv. data_shape = (2, 2, 4, 3) # NHWC data_dtype = "uint8" kernel_shape = (2, 2, 3, 1) # HWOI kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=5, kernel_zero_point=3, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(1, 1), groups=3, data_layout="NHWC", kernel_layout="HWOI", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype)
def test_dilation(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # Non-zero kernel point - fall back to simpler lowering. data_shape = (2, 4, 4, 4) data_dtype = "uint8" kernel_shape = (3, 4, 2, 2) kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=5, kernel_zero_point=3, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(2, 2), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype) # Zero kernel point data_shape = (2, 4, 4, 4) data_dtype = "uint8" kernel_shape = (3, 4, 2, 2) kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=0, kernel_zero_point=0, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 1), dilation=(2, 2), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32", ) verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype)
def test_qnn_legalize(): """Test directly replacing an operator with a new one""" def before(): x = relay.var("x", shape=(1, 64, 56, 56), dtype='int8') y = relay.qnn.op.requantize(x, input_scale=1, input_zero_point=0, output_scale=1, output_zero_point=0, out_dtype='int8') y = relay.Function([x], y) return y def legalize_qnn_requantize(attrs, inputs, types): data = inputs[0] data = relay.add(relay.const(0, 'int8'), data) y = relay.qnn.op.requantize(data, input_scale=1, input_zero_point=0, output_scale=1, output_zero_point=0, out_dtype='int8') return y def expected(): x = relay.var("x", shape=(1, 64, 56, 56), dtype='int8') y = relay.add(relay.const(0, 'int8'), x) z = relay.qnn.op.requantize(y, input_scale=1, input_zero_point=0, output_scale=1, output_zero_point=0, out_dtype='int8') z = relay.Function([x], z) return z a = before() with TempOpAttr("qnn.requantize", "FTVMQnnLegalize", legalize_qnn_requantize): # Check that Relay Legalize does not change the graph. a = run_opt_pass(a, relay.transform.Legalize()) b = run_opt_pass(before(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a) # Check that QNN Legalize modifies the graph. a = run_opt_pass(a, relay.qnn.transform.Legalize()) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_select_implementation(): target = tvm.target.create("llvm") def _select_impl(dshape, wshape, use_autotvm=False): data = relay.var("data", shape=dshape) weight = relay.var("wshape", shape=wshape) out = relay.nn.conv2d(data, weight, padding=(1, 1)) out = run_infer_type(out) return relay.backend.compile_engine.select_implementation( relay.op.get("nn.conv2d"), out.attrs, [tvm.placeholder(dshape), tvm.placeholder(wshape)], out.checked_type, target, use_autotvm) with TempOpAttr("nn.conv2d", "FTVMStrategy", _tmp_strategy): impl, _ = _select_impl((1, 8, 7, 7), (32, 8, 3, 3)) assert impl.name == "conv2d_2" impl, _ = _select_impl((1, 8, 7, 7), (32, 8, 3, 3), True) assert impl.name == "conv2d_2" impl, _ = _select_impl((1, 16, 7, 7), (32, 16, 3, 3)) assert impl.name == "conv2d_3" impl, _ = _select_impl((1, 16, 7, 7), (32, 16, 3, 3), True) assert impl.name == "conv2d_3" # add autotvm record records = [] records.append( _create_record("test/conv2d_1", (1, 8, 7, 7), (32, 8, 3, 3), target, 0.5)) records.append( _create_record("test/conv2d_1", (1, 16, 7, 7), (32, 16, 3, 3), target, 1.0)) with target: with autotvm.apply_history_best(records): impl, _ = _select_impl((1, 8, 7, 7), (32, 8, 3, 3), True) assert impl.name == "conv2d_1" impl, _ = _select_impl((1, 16, 7, 7), (32, 16, 3, 3), True) assert impl.name == "conv2d_1" records.append( _create_record("test/conv2d_2", (1, 8, 7, 7), (32, 8, 3, 3), target, 0.2)) records.append( _create_record("test/conv2d_1", (1, 16, 7, 7), (32, 16, 3, 3), target, 1.2)) with target: with autotvm.apply_history_best(records): impl, _ = _select_impl((1, 8, 7, 7), (32, 8, 3, 3), True) assert impl.name == "conv2d_2" impl, _ = _select_impl((1, 16, 7, 7), (32, 16, 3, 3), True) assert impl.name == "conv2d_1"
def test_alter_layout_lrn(): """Test alternating the layout of a conv2d. The layout of broadcast operators and the weight should be changed accordingly. """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias") weight = relay.var("weight") y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.max_pool2d(y, pool_size=(2, 2)) y = relay.nn.lrn(y) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs["data_layout"] = "NCHW16c" new_attrs["kernel_layout"] = "OIHW16i" return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias", shape=(64,)) weight = relay.var("weight", shape=(64, 64, 3, 3)) y = relay.layout_transform(x, "NCHW", "NCHW16c") w = relay.layout_transform(weight, "OIHW", "OIHW16i") y = relay.nn.conv2d( y, w, channels=64, kernel_size=(3, 3), padding=(1, 1), kernel_layout="OIHW16i", data_layout="NCHW16c", ) y = relay.nn.max_pool2d(y, pool_size=(2, 2), layout="NCHW16c") y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.nn.lrn(y) y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_padding(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (1, 4, 2, 2) data_dtype = 'uint8' kernel_shape = (3, 4, 2, 2) kernel_dtype = 'uint8' ref_func, qnn_func = get_funcs(data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=8, kernel_zero_point=5, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(1, 1), strides=(1, 1), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32") verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype) # Try different layout data_shape = (2, 2, 4, 4) # NHWC data_dtype = 'uint8' kernel_shape = (2, 2, 4, 3) # HWIO kernel_dtype = 'uint8' ref_func, qnn_func = get_funcs(data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=8, kernel_zero_point=3, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(1, 1), strides=(1, 1), dilation=(1, 1), data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32") verify(ref_func, qnn_func, data_shape, data_dtype, kernel_shape, kernel_dtype)
def _test_legalize_dense(data_shape, kernel_shape, pad_shape, do_pad=True): """test legalize dense to enable tensorcore""" M, K = data_shape N, _ = kernel_shape out_shape = (M, N) dm, dk, dn = pad_shape def before(): x = relay.var("x", shape=data_shape, dtype="float16") weight = relay.var("weight", shape=kernel_shape, dtype="float16") y = relay.nn.dense(x, weight) y = relay.Function([x, weight], y) return y def legalize_dense(attrs, inputs, types): with tvm.target.Target("cuda"): return topi.nn.dense_legalize(attrs, inputs, types) def expected(): if not do_pad: return before() x = relay.var("x", shape=data_shape, dtype="float16") if dm or dk: x_pad = relay.nn.pad(x, pad_width=((0, dm), (0, dk))) else: x_pad = x weight = relay.var("weight", shape=(kernel_shape), dtype="float16") if dn or dk: weight_pad = relay.nn.pad(weight, pad_width=((0, dn), (0, dk))) else: weight_pad = weight y_pad = relay.nn.dense( x_pad, weight_pad, ) if dm or dn: y = relay.strided_slice(y_pad, begin=[0, 0], end=out_shape) else: y = y_pad y = relay.Function([x, weight], y) return y with TempOpAttr("nn.dense", "FTVMLegalize", legalize_dense): a = before() a = run_opt_pass(a, transform.Legalize()) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal( a, b), "Actual = \n" + str(a) + "Expected = \n" + str(b)
def test_alter_layout_scalar(): """Test alternating the layout of a conv2d. The layout of broadcast operators and the weight should be changed accordingly. """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var("weight") y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.add(y, relay.const(1, "float32")) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs['data_layout'] = 'NCHW16c' return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) w = relay.var("weight") y = relay.layout_transform(x, "NCHW", "NCHW16c") y = relay.nn.conv2d(y, w, channels=64, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c") y = relay.add(y, relay.const(1.0, "float32")) y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass( a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_tflite_output_multiplier_greater_than_one(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (2, 1, 2, 4) data_dtype = 'uint8' kernel_shape = (3, 1, 2, 2) kernel_dtype = 'uint8' ref_func, qnn_func = get_funcs(data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_scale=1.0, kernel_scale=1.0, input_zero_point=128, kernel_zero_point=128, kernel_size=(2, 2), padding=(0, 0), strides=(2, 2), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32") golden_data = 128 + np.array((1, 1, 1, 1, 2, 2, 2, 2, 1, 2, 3, 4, 1, 2, 3, 4)).reshape(data_shape).astype('uint8') golden_weight = 128 + np.array((1, 2, 3, 4, -1, 1, -1, 1, -1, -1, 1, 1)).reshape(kernel_shape) golden_weight = golden_weight.astype('uint8') with relay.build_config(opt_level=2): params = {'kernel': golden_weight} graph, lib, params = relay.build(qnn_func, "llvm", params=params) mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0)) mod.set_input("data", golden_data) mod.set_input(**params) mod.run() qnn_output = mod.get_output(0).asnumpy() golden_output = np.array((17, 17, 0, 0, 2, 2, 16, 36, 2, 2, 0, 0)).reshape(2, 3, 1, 2) np.testing.assert_equal(qnn_output, golden_output)
def test_alter_layout_nchw_upsamping_op(): """Test upsamping operators """ def before(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var('weight', shape=(32, 32, 3, 3)) y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.upsampling(y, scale_h=2, scale_w=2) y = relay.nn.avg_pool2d(y, pool_size=(2, 2), strides=(2, 2)) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs['data_layout'] = 'NCHW16c' return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var("weight") x = relay.layout_transform(x, "NCHW", "NCHW16c") y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c") y = relay.nn.upsampling(y, scale_h=2, scale_w=2, layout="NCHW16c") y = relay.nn.avg_pool2d(y, pool_size=(2, 2), strides=(2, 2), layout='NCHW16c') y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_alter_layout_depthwise_conv2d(): """Test depthwise_conv2d operator""" def before(): x = relay.var("x", shape=(1, 32, 56, 56)) w = relay.var("w", shape=(32, 1, 3, 3)) y = relay.nn.conv2d(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), groups=32) y = relay.Function(analysis.free_vars(y), y) return y import topi def alter_conv2d(attrs, inputs, tinfos, out_type): with tvm.target.create("llvm"): return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, out_type) def expected(): x = relay.var("x", shape=(1, 32, 56, 56)) w = relay.var("w", shape=(32, 1, 3, 3)) x = relay.layout_transform(x, "NCHW", "NCHW8c") w = relay.layout_transform(w, "OIHW", "OIHW1i8o") y = relay.nn.contrib_depthwise_conv2d_nchwc(x, w, padding=(1, 1, 1, 1), channels=32, kernel_size=(3, 3), groups=32, data_layout="NCHW8c", kernel_layout="OIHW1i8o", out_layout="NCHW8c") y = relay.layout_transform(y, "NCHW8c", "NCHW") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass( a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert (analysis.alpha_equal(a, b))
def test_alter_layout_prelu(): """Test PRelu operator""" def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var("weight") alpha = relay.var("alpha", relay.IncompleteType()) y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.prelu(y, alpha) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs['data_layout'] = 'NCHW16c' return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) w = relay.var("weight") alpha = relay.var("alpha", relay.IncompleteType()) y = relay.layout_transform(x, "NCHW", "NCHW16c") y = relay.nn.conv2d(y, w, channels=64, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c") y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.nn.prelu(y, alpha) y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass( a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert (analysis.alpha_equal(a, b))
def test_alter_layout_broadcast_op(): """Test boradcast operators """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias", shape=(64,)) scale = relay.var("scale", shape=(64, 1, 1)) weight = relay.var("weight") y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.bias_add(y, bias) # test broadcasting to lhs y = relay.multiply(scale, y) # test broadcasting to rhs y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs["data_layout"] = "NCHW16c" return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias", shape=(64,)) scale = relay.var("scale", shape=(64, 1, 1)) weight = relay.var("weight") x = relay.layout_transform(x, "NCHW", "NCHW16c") bias = relay.expand_dims(bias, 1, 2) bias = relay.expand_dims(bias, 0, 1) bias = relay.layout_transform(bias, "NCHW", "NCHW16c") scale = relay.expand_dims(scale, 0, 1) scale = relay.layout_transform(scale, "NCHW", "NCHW16c") y = relay.nn.conv2d( x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c" ) y = relay.add(y, bias) # test broadcasting to lhs y = relay.multiply(scale, y) # test broadcasting to rhs y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_alter_layout_resnet(): """Test alternating the layout of a residual block This also tests the elimination of duplicated transformation. If a same transformation applies to a same node twice, only one transformation will be created. """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1") weight2 = relay.var("weight2") y = relay.nn.conv2d(x, weight1, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) y2 = relay.nn.conv2d(x, weight2, channels=32, kernel_size=(1, 1)) y2 = relay.nn.relu(y2) y = y + y2 y = relay.nn.global_max_pool2d(y) return relay.Function(analysis.free_vars(y), y) def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs["data_layout"] = "NCHW16c" return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1") weight2 = relay.var("weight2") x = relay.layout_transform(x, "NCHW", "NCHW16c") y = relay.nn.conv2d( x, weight1, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c" ) y = relay.nn.relu(y) y2 = relay.nn.conv2d(x, weight2, channels=32, kernel_size=(1, 1), data_layout="NCHW16c") y2 = relay.nn.relu(y2) y = y + y2 y = relay.nn.global_max_pool2d(y, layout="NCHW16c") y = relay.layout_transform(y, "NCHW16c", "NCHW") return relay.Function(analysis.free_vars(y), y) with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_op_temp_attr(): """Tests reset_attr functionality.""" def add1(x): return x + 1 def add2(x): return x + 2 # Set original attr value is add1. tvm.ir.register_op_attr("sqrt", "ftest", add1) with TempOpAttr("sqrt", "ftest", add2): # Check that the attr value is updated to add2. assert relay.op.get("sqrt").get_attr("ftest")(1) == 3 # Check that the attr value is recovered to add1. assert relay.op.get("sqrt").get_attr("ftest")(1) == 2
def test_alter_op_with_global_var(): """Test directly replacing an operator with a new one""" def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var('weight', shape=(64, 64, 3, 3)) y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) mod = tvm.IRModule() foo = relay.GlobalVar('foo') mod[foo] = relay.Function([x, weight], y) mod["main"] = relay.Function([x, weight], foo(x, weight)) return mod def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs weight = relay.multiply(weight, relay.const(2.0, "float32")) return relay.nn.conv2d(data, weight, **attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var('weight', shape=(64, 64, 3, 3)) y = relay.nn.conv2d(x, relay.multiply(weight, relay.const(2.0, "float32")), channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) mod = tvm.IRModule() foo = relay.GlobalVar('foo') mod[foo] = relay.Function([x, weight], y) mod["main"] = relay.Function([x, weight], foo(x, weight)) return mod with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = transform.AlterOpLayout()(a) b = transform.InferType()(expected()) assert tvm.ir.structural_equal(a, b, map_free_vars=True), "Actual = \n" + str(a)
def test_get_valid_implementations(): target = tvm.target.create("llvm") def _get_impls(dshape, wshape): data = relay.var("data", shape=dshape) weight = relay.var("wshape", shape=wshape) out = relay.nn.conv2d(data, weight, padding=(1, 1)) out = run_infer_type(out) return relay.backend.compile_engine.get_valid_implementations( relay.op.get("nn.conv2d"), out.attrs, [tvm.placeholder(dshape), tvm.placeholder(wshape)], out.checked_type, target) with TempOpAttr("nn.conv2d", "FTVMStrategy", _tmp_strategy): impls = _get_impls((1, 8, 7, 7), (32, 8, 3, 3)) assert len(impls) == 2 impls = _get_impls((1, 16, 7, 7), (32, 16, 3, 3)) assert len(impls) == 3
def test_alter_layout_strided_slice(): """Test rewriting strided_slice during alter_iop_layout""" def before(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var('weight', shape=(32, 32, 3, 3)) y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.strided_slice(y, begin=[0, 16], end=[None, None]) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs['data_layout'] = 'NCHW4c' return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var("weight") x = relay.layout_transform(x, "NCHW", "NCHW4c") y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW4c") y = relay.strided_slice(y, begin=[0, 4], end=[None, 8]) y = relay.layout_transform(y, "NCHW4c", "NCHW") y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass( a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_legalize_none(): """Test doing nothing by returning 'None' """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) y = relay.nn.global_max_pool2d(x) y = relay.Function([x], y) return y called = [False] def legalize_conv2d(attrs, inputs, types): called[0] = True return None with TempOpAttr("nn.global_max_pool2d", "FTVMLegalize", legalize_conv2d): a = before() a = run_opt_pass(a, transform.Legalize()) b = run_opt_pass(before(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) assert called[0]
def test_alter_return_none(): """Test doing nothing by returning 'None' """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) y = relay.nn.global_max_pool2d(x) y = relay.Function([x], y) return y called = [False] def alter_conv2d(attrs, inputs, tinfos, out_type): called[0] = True return None with TempOpAttr("nn.global_max_pool2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(before(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a) assert (called[0])
def test_tflite_anistropic_strides(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (1, 1, 3, 6) data_dtype = 'uint8' kernel_shape = (1, 1, 2, 2) kernel_dtype = 'uint8' ref_func, qnn_func = get_funcs(data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=127, kernel_zero_point=127, input_scale=1.0, kernel_scale=1.0, kernel_size=(2, 2), padding=(0, 0), strides=(1, 3), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32") golden_data = np.array( (133, 131, 129, 125, 123, 121, 135, 133, 131, 123, 121, 119, 137, 135, 133, 121, 119, 117)).reshape(data_shape) golden_data = golden_data.astype('uint8') golden_weight = np.array((129, 131, 133, 135)).reshape(kernel_shape) golden_weight = golden_weight.astype('uint8') with relay.build_config(opt_level=2): params = {'kernel': golden_weight} graph, lib, params = relay.build(qnn_func, "llvm", params=params) mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0)) mod.set_input("data", golden_data) mod.set_input(**params) mod.run() qnn_output = mod.get_output(0).asnumpy() golden_output = np.array((124, -92, 164, -132)).reshape(1, 1, 2, 2) np.testing.assert_equal(qnn_output, golden_output)
def test_legalize(): """Test directly replacing an operator with a new one""" def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var("weight", shape=(64, 64, 3, 3)) y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) y = relay.Function([x, weight], y) return y def legalize_conv2d(attrs, inputs, types): data, weight = inputs weight = relay.multiply(weight, relay.const(2.0, "float32")) return relay.nn.conv2d(data, weight, **attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var("weight", shape=(64, 64, 3, 3)) y = relay.nn.conv2d( x, relay.multiply(weight, relay.const(2.0, "float32")), channels=64, kernel_size=(3, 3), padding=(1, 1), ) y = relay.nn.relu(y) y = relay.Function([x, weight], y) return y with TempOpAttr("nn.conv2d", "FTVMLegalize", legalize_conv2d): a = before() a = run_opt_pass(a, transform.Legalize()) b = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_alter_layout_nhwc_arm(): """ Check that AlterOplayout does not alter NHWC data layout. """ def alter_conv2d(attrs, inputs, tinfos, out_type): import topi with tvm.target.create("llvm -device=arm_cpu"): return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, out_type) # Check NHWC conversion. def before_nhwc(): x = relay.var("x", shape=(1, 56, 56, 64)) weight1 = relay.var('weight1', shape=(3, 3, 64, 64)) weight2 = relay.var('weight2', shape=(3, 3, 64, 64)) y = relay.nn.conv2d(x, weight1, channels=64, kernel_size=(3, 3), data_layout='NHWC', kernel_layout='HWIO') y = relay.nn.relu(y) y = relay.nn.avg_pool2d(y, pool_size=(1, 1), layout='NHWC') y = relay.nn.conv2d(y, weight2, channels=64, kernel_size=(3, 3), data_layout='NHWC', kernel_layout='HWIO') y = relay.nn.relu(y) y = relay.Function(analysis.free_vars(y), y) return y def expected_nhwc(): return before_nhwc() with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before_nhwc() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected_nhwc(), transform.InferType()) assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
def test_tflite_large_irregular(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # uint8 input data_shape = (1, 1024, 1, 1) data_dtype = "uint8" kernel_shape = (1001, 1024, 1, 1) kernel_dtype = "uint8" ref_func, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=127, kernel_zero_point=127, input_scale=1.0, kernel_scale=1.0, kernel_size=(1, 1), padding=(0, 0), strides=(1, 1), dilation=(1, 1), data_layout="NCHW", kernel_layout="OIHW", out_dtype="int32", ) golden_data = np.full(data_shape, 127).astype("uint8") golden_weight = np.full(kernel_shape, 127).astype("uint8") with tvm.transform.PassContext(opt_level=2): params = {"kernel": golden_weight} graph, lib, params = relay.build(qnn_func, "llvm", params=params) mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0)) mod.set_input("data", golden_data) mod.set_input(**params) mod.run() qnn_output = mod.get_output(0).asnumpy() golden_output = np.full((1, 1001, 1, 1), 0).astype("uint8") np.testing.assert_equal(qnn_output, golden_output)
def test_broadcast_layout(): with TempOpAttr("qnn.conv2d", "FTVMQnnLegalize", legalize_qnn_conv2d): # Test broadcast support for NHWC layout. data_shape = (1, 229, 229, 3) # NHWC data_dtype = "uint8" kernel_shape = (7, 7, 3, 64) # HWIO kernel_dtype = "int8" _, qnn_func = get_funcs( data_shape=data_shape, data_dtype=data_dtype, kernel_shape=kernel_shape, kernel_dtype=kernel_dtype, input_zero_point=8, kernel_zero_point=3, input_scale=1.0, kernel_scale=1.0, kernel_size=(7, 7), padding=(1, 1), strides=(1, 1), dilation=(1, 1), data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32", ) func = qnn_func["main"].body bias = relay.var("bias", shape=(64,), dtype="int32") bias2 = relay.var("bias2", shape=(1, 225, 225, 1), dtype="int32") # Check broadcast support on both lhs and rhs func = relay.add(func, bias2) func = relay.add(bias2, func) func = relay.add(bias, func) func = relay.add(func, bias) func = relay.Function(relay.analysis.free_vars(func), func) mod = tvm.IRModule.from_expr(func) with tvm.transform.PassContext(opt_level=3): graph, lib, params = relay.build(mod, "llvm -mcpu=skylake-avx512")
def test_alter_op(): """Test directly replacing an operator with a new one""" def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var('weight', shape=(64, 64, 3, 3)) y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) y = relay.Function([x, weight], y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs weight = relay.multiply(weight, relay.const(2.0, "float32")) return relay.nn.conv2d(data, weight, **attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight = relay.var('weight', shape=(64, 64, 3, 3)) y = relay.nn.conv2d(x, relay.multiply(weight, relay.const(2.0, "float32")), channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.relu(y) y = relay.Function([x, weight], y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass(a, transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)
def test_alter_layout(): """Test alternating the layout of a conv2d. The layout of broadcast operators and the weight should be changed accordingly. """ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias") weight = relay.var("weight") y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) y = relay.nn.bias_add(y, bias) # a useless tuple, which will be eliminated y = relay.Tuple([y])[0] y = relay.nn.relu(y) y = relay.nn.max_pool2d(y, pool_size=(2, 2)) y = relay.cast(y, 'int32') y = relay.nn.batch_flatten(y) y = relay.Function(analysis.free_vars(y), y) return y def alter_conv2d(attrs, inputs, tinfos, out_type): data, weight = inputs new_attrs = dict(attrs) new_attrs['data_layout'] = 'NCHW16c' new_attrs['kernel_layout'] = 'OIHW16i' return relay.nn.conv2d(data, weight, **new_attrs) def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) bias = relay.var("bias", shape=(64, )) weight = relay.var("weight", shape=(64, 64, 3, 3)) y = relay.layout_transform(x, "NCHW", "NCHW16c") w = relay.layout_transform(weight, "OIHW", "OIHW16i") y = relay.nn.conv2d(y, w, channels=64, kernel_size=(3, 3), padding=(1, 1), kernel_layout="OIHW16i", data_layout="NCHW16c") b = relay.expand_dims(bias, axis=1, num_newaxis=2) b = relay.expand_dims(b, axis=0, num_newaxis=1) b = relay.layout_transform(b, "NCHW", "NCHW16c") y = relay.add(y, b) y = relay.nn.relu(y) y = relay.nn.max_pool2d(y, pool_size=(2, 2), layout="NCHW16c") y = relay.cast(y, 'int32') y = relay.layout_transform(y, "NCHW16c", "NCHW") y = relay.nn.batch_flatten(y) y = relay.Function(analysis.free_vars(y), y) return y with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = before() a = run_opt_pass( a, [transform.CanonicalizeOps(), transform.AlterOpLayout()]) b = run_opt_pass(expected(), transform.InferType()) assert analysis.alpha_equal(a, b), "Actual = \n" + str(a)