From bc834de395d8e6b78b432a013e15998997522d0c Mon Sep 17 00:00:00 2001 From: Animesh Jain Date: Thu, 3 Oct 2019 00:58:52 +0000 Subject: [PATCH] [Relay][CanonicalizeOps] Makeing Bias_add shape same as the other operand. --- src/relay/pass/pattern_util.h | 16 +++++++-- tests/python/relay/test_op_qnn_conv2d.py | 33 +++++++++++++++++++ .../python/relay/test_pass_alter_op_layout.py | 6 ++-- .../python/relay/test_pass_fold_scale_axis.py | 28 +++++++++------- .../relay/test_pass_simplify_inference.py | 3 ++ 5 files changed, 70 insertions(+), 16 deletions(-) diff --git a/src/relay/pass/pattern_util.h b/src/relay/pass/pattern_util.h index bf9621bb404a..dfa99e201ba2 100644 --- a/src/relay/pass/pattern_util.h +++ b/src/relay/pass/pattern_util.h @@ -148,13 +148,23 @@ inline Expr ExpandBiasToMatchAxis(Expr bias, static const Op& expand_dims = Op::Get("expand_dims"); for (size_t i = axes.size(); i != 0; --i) { if (i == axes.size()) { - int64_t num_pad_axis = target_ndim - axes[i - 1]->value - 1; - if (num_pad_axis > 0) { + // The idea is to make the shape of the operand same as target_ndim. + // Therefore, there are 2 expand dims, one before the axis and one after + // the axis. + int64_t num_pad_axis_inner = target_ndim - axes[i - 1]->value - 1; + int64_t num_pad_axis_outer = target_ndim - num_pad_axis_inner - 1; + if (num_pad_axis_inner > 0) { auto attrs = make_node(); attrs->axis = i; - attrs->num_newaxis = static_cast(num_pad_axis); + attrs->num_newaxis = static_cast(num_pad_axis_inner); bias = CallNode::make(expand_dims, {bias}, Attrs(attrs), {}); } + if (num_pad_axis_outer > 0) { + auto new_attrs = make_node(); + new_attrs->axis = 0; + new_attrs->num_newaxis = static_cast(num_pad_axis_outer); + bias = CallNode::make(expand_dims, {bias}, Attrs(new_attrs), {}); + } } else { int64_t diff = axes[i]->value - axes[i - 1]->value; CHECK_GE(diff, 0L); diff --git a/tests/python/relay/test_op_qnn_conv2d.py b/tests/python/relay/test_op_qnn_conv2d.py index c8e479d99ee4..90f0eafdbc1a 100644 --- a/tests/python/relay/test_op_qnn_conv2d.py +++ b/tests/python/relay/test_op_qnn_conv2d.py @@ -607,6 +607,38 @@ def tflite_anistropic_strides(): golden_output = np.array((124, -92, 164, -132)).reshape(1, 1, 2, 2) np.testing.assert_equal(qnn_output, golden_output) +def broadcast_layout_test(): + # 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, + 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") + + # Check broadcast support on both lhs and rhs + func = relay.nn.bias_add(func, bias, axis=3) + # func = relay.nn.bias_add(bias, func) + + func = relay.Function(relay.analysis.free_vars(func), func) + mod = relay.Module.from_expr(func) + mod = transform.CanonicalizeOps()(mod) + with relay.build_config(opt_level=3): + graph, lib, params = relay.build(mod, "llvm -mcpu=skylake-avx512") + if __name__ == "__main__": no_zero_point_test() input_zero_point_test() @@ -620,3 +652,4 @@ def tflite_anistropic_strides(): tflite_large_irregular_test() tflite_output_multiplier_greater_than_one() tflite_anistropic_strides() + broadcast_layout_test() diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index 6b31eed8f166..13481ad14471 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -134,7 +134,8 @@ def expected(): kernel_layout="OIHW16i", data_layout="NCHW16c") b = relay.expand_dims(bias, axis=1, num_newaxis=2) - b = relay.layout_transform(b, "CHW", "CHW16c") + 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) @@ -304,7 +305,8 @@ def expected(): weight = relay.var("weight") x = relay.layout_transform(x, "NCHW", "NCHW16c") bias = relay.expand_dims(bias, 1, 2) - bias = relay.layout_transform(bias, "CHW", "CHW16c") + bias = relay.expand_dims(bias, 0, 1) + bias = relay.layout_transform(bias, "NCHW", "NCHW16c") scale = relay.layout_transform(scale, "CHW", "CHW16c") y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW16c") diff --git a/tests/python/relay/test_pass_fold_scale_axis.py b/tests/python/relay/test_pass_fold_scale_axis.py index d6f471bef04a..7a5a5e427c3f 100644 --- a/tests/python/relay/test_pass_fold_scale_axis.py +++ b/tests/python/relay/test_pass_fold_scale_axis.py @@ -51,11 +51,12 @@ def expected(x, conv_weight, in_bias, in_scale, channels): args = [x, conv_weight, in_bias] in_bias = relay.expand_dims(in_bias, axis=1, num_newaxis=2) squeezed_scale = relay.squeeze(in_scale, axis=[1,2]) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=0, num_newaxis=1) x = relay.nn.relu(x) - in_bias = relay.divide(in_bias, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2)) + in_bias = relay.divide(in_bias, squeezed_scale) x = relay.add(x, in_bias) - conv_weight = relay.multiply( - conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2)) + conv_weight = relay.multiply(conv_weight, squeezed_scale) y = relay.nn.conv2d(x, conv_weight, channels=channels, kernel_size=(3, 3), @@ -109,10 +110,13 @@ def before(x, conv_weight, in_bias, in_scale, channels): def expected(x, conv_weight, in_bias, in_scale, channels): args = [x, conv_weight, in_bias] x = relay.nn.relu(x) - in_bias = relay.divide(in_bias, in_scale) + in_scale1 = relay.expand_dims(in_scale, 0, 3) + in_scale2 = relay.expand_dims(in_scale, 0, 3) + in_scale3 = relay.expand_dims(in_scale, 0, 3) + in_bias = relay.divide(in_bias, in_scale1) x = relay.subtract(x, in_bias) y1 = relay.nn.conv2d(x, - relay.multiply(conv_weight, in_scale), + relay.multiply(conv_weight, in_scale2), channels=channels, kernel_size=(3, 3), data_layout="NHWC", @@ -120,7 +124,7 @@ def expected(x, conv_weight, in_bias, in_scale, channels): groups=channels, padding=(1, 1)) y2 = relay.nn.conv2d(x, - relay.multiply(conv_weight, in_scale), + relay.multiply(conv_weight, in_scale3), channels=channels, kernel_size=(3, 3), data_layout="NHWC", @@ -225,8 +229,9 @@ def expected(x, conv_weight, in_scale, channels): # use a fixed order of args so alpha equal check can pass args = [x, conv_weight] squeezed_scale = relay.squeeze(in_scale, axis=[1,2]) - conv_weight = relay.multiply( - conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2)) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=0, num_newaxis=1) + conv_weight = relay.multiply(conv_weight, squeezed_scale) y = relay.nn.conv2d(x, conv_weight, channels=channels, @@ -271,14 +276,15 @@ def expected(x, conv_weight, out_bias, out_scale, channels): out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2) squeezed_scale = relay.squeeze(out_scale, axis=[1,2]) conv_weight = relay.multiply( - conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3)) + conv_weight, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3)) y = relay.nn.conv2d(x, conv_weight, channels=channels, kernel_size=(3, 3), padding=(1, 1)) - out_bias = relay.multiply(out_bias, - relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2)) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2) + squeezed_scale = relay.expand_dims(squeezed_scale, axis=0, num_newaxis=1) + out_bias = relay.multiply(out_bias, squeezed_scale) y = relay.add(y, out_bias) y = relay.nn.relu(y) return relay.Function(args, y) diff --git a/tests/python/relay/test_pass_simplify_inference.py b/tests/python/relay/test_pass_simplify_inference.py index 4e62fa6dcb08..969f07a664bd 100644 --- a/tests/python/relay/test_pass_simplify_inference.py +++ b/tests/python/relay/test_pass_simplify_inference.py @@ -29,6 +29,9 @@ def simple_bn(x, gamma, beta, moving_mean, moving_var, if num_newaxis: scale = rly.expand_dims(scale, axis=1, num_newaxis=num_newaxis) shift = rly.expand_dims(shift, axis=1, num_newaxis=num_newaxis) + if num_newaxis != 3: + shift = rly.expand_dims(shift, axis=0, num_newaxis=1) + scale = rly.expand_dims(scale, axis=0, num_newaxis=1) return x * scale + shift def check(dim, axis, nstep):