From 52a87d40d2cd2d6ac63b4f9e9dc695c5b5cae2d2 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 17:02:53 -0800 Subject: [PATCH 1/7] Add compute, schedule, and tests for expand_dims and squeeze --- python/tvm/relay/op/_transform.py | 44 +++++++++++++++++++++++++++- tests/python/relay/test_op_level1.py | 16 ++++++++++ tests/python/relay/test_op_level3.py | 15 ++++++++++ 3 files changed, 74 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 01814e0f73e0..14763b5e2a56 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -1,8 +1,50 @@ #pylint: disable=invalid-name, unused-argument """Backend compiler related feature registration""" from __future__ import absolute_import +import topi +import topi.cuda +from tvm import container from . import op as _reg -from .op import schedule_injective, OpPattern +from .op import (schedule_broadcast, schedule_injective, + register_compute, register_schedule, + register_pattern, OpPattern) + +# squeeze +@register_compute("squeeze") +def squeeze_compiler(attrs, inputs, output_type, target): + """Compiler for squeeze dims.""" + assert len(inputs) == 1 + + if attrs.axis is None: + axis = None + elif isinstance(attrs.axis, container.Array): + axis = tuple(attrs.axis) + else: + axis = int(attrs.axis) + + return [topi.squeeze(inputs[0], axis)] + +register_pattern("squeeze", OpPattern.INJECTIVE) +register_schedule("squeeze", schedule_injective) + +# expand_dims +@register_compute("expand_dims") +def expand_dims_compiler(attrs, inputs, output_type, target): + """Compiler for expand_dims.""" + assert len(inputs) == 1 + + new_axis = int(attrs.num_newaxis) + assert new_axis >= 0 + + # axis should be in range [-data.ndim - 1, data.ndim] + axis = int(attrs.axis) + assert axis >= -len(inputs[0].shape) - 1 + assert axis <= len(inputs[0].shape) + + return [topi.expand_dims(inputs[0], axis, new_axis)] + +_reg.register_schedule("expand_dims", schedule_broadcast) +_reg.register_pattern("expand_dims", OpPattern.BROADCAST) # strided_slice _reg.register_schedule("strided_slice", schedule_injective) diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index 35844ddd4a3f..7e570862d230 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -90,6 +90,22 @@ def check_binary_op(opfunc, ref): check_binary_op(opfunc, ref) +def test_expand_dims(): + # based on topi test + def verify_expand_dims(dshape, dtype, oshape, axis, num_newaxis): + x = relay.Var("x", relay.TensorType(dshape, dtype)) + func = relay.Function([x], relay.expand_dims(x, axis, num_newaxis)) + for target, ctx in ctx_list(): + data = np.random.uniform(size=dshape).astype(dtype) + ref_res = data.reshape(oshape) + intrp = relay.create_executor("graph", ctx=ctx, target=target) + op_res = intrp.evaluate(func)(data) + np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01) + + verify_expand_dims((3, 10), 'float32', (3, 10, 1, 1), 2, 2) + verify_expand_dims((3, 10), 'float32', (1, 3, 10), -3, 1) + + def test_bias_add(): xshape=(10, 2, 3, 4) bshape=(2,) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 22469cc7fdbe..58315e31f553 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -60,6 +60,21 @@ def test_clip(): np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01) +def test_squeeze(): + def verify_squeeze(shape, dtype, axis): + x = relay.var("x", relay.TensorType(shape, dtype)) + squeeze = relay.squeeze(x, axis=axis) + + data = np.random.random_sample(shape).astype(dtype) + intrp = create_executor() + op_res = intrp.evaluate(squeeze, { x : relay.const(data) }) + ref_res = np.squeeze(data, axis=axis) + np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01) + + verify_squeeze((1, 3, 2, 5), "float32", None) + verify_squeeze((1, 3, 1), "float32", 0) + verify_squeeze((1, 3, 2, 5, 1), "float32", -1) + verify_squeeze((1, 2, 1, 2, 1), "float32", (0, 2, 4)) def test_transpose_infer_type(): From c500be5341d85d999c3650fd701e2077f41bec31 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 17:25:13 -0800 Subject: [PATCH 2/7] Fix import --- python/tvm/relay/op/_transform.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 14763b5e2a56..62a9e0b6501c 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -5,9 +5,9 @@ import topi.cuda from tvm import container from . import op as _reg -from .op import (schedule_broadcast, schedule_injective, - register_compute, register_schedule, +from .op import (schedule_injective, register_compute, register_schedule, register_pattern, OpPattern) +from .tensor import schedule_broadcast # squeeze @register_compute("squeeze") From 9104d414d38b62408b035fca3d5a0432e7321788 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 17:35:07 -0800 Subject: [PATCH 3/7] Really fix import --- python/tvm/relay/op/_transform.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 62a9e0b6501c..cd32aea38d90 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -7,7 +7,8 @@ from . import op as _reg from .op import (schedule_injective, register_compute, register_schedule, register_pattern, OpPattern) -from .tensor import schedule_broadcast + +schedule_broadcast = schedule_injective # squeeze @register_compute("squeeze") From 94672f05d8d694b11848a56bb2a79e73d65f89ba Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 18:10:02 -0800 Subject: [PATCH 4/7] Fix squeeze test --- tests/python/relay/test_op_level3.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 58315e31f553..a68dd35ebe03 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -72,9 +72,9 @@ def verify_squeeze(shape, dtype, axis): np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01) verify_squeeze((1, 3, 2, 5), "float32", None) - verify_squeeze((1, 3, 1), "float32", 0) - verify_squeeze((1, 3, 2, 5, 1), "float32", -1) - verify_squeeze((1, 2, 1, 2, 1), "float32", (0, 2, 4)) + verify_squeeze((1, 3, 1), "float32", [0]) + verify_squeeze((1, 3, 2, 5, 1), "float32", [-1]) + verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) def test_transpose_infer_type(): @@ -310,6 +310,7 @@ def test_infer_type_prelu(): test_full_like() test_infer_type_leaky_relu() test_infer_type_prelu() + test_squeeze() test_squeeze_infer_type() test_squeeze_bad_axes_infer_type() test_split_infer_type() From b719e5008cbb986afc838eb27b40ea9c5b5b202f Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 18:21:25 -0800 Subject: [PATCH 5/7] Have to convert lists to tuples for numpy --- tests/python/relay/test_op_level3.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index a68dd35ebe03..2862cfe10ac1 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -65,10 +65,12 @@ def verify_squeeze(shape, dtype, axis): x = relay.var("x", relay.TensorType(shape, dtype)) squeeze = relay.squeeze(x, axis=axis) + np_axis = tuple(axis) if axis is not None else None + data = np.random.random_sample(shape).astype(dtype) intrp = create_executor() op_res = intrp.evaluate(squeeze, { x : relay.const(data) }) - ref_res = np.squeeze(data, axis=axis) + ref_res = np.squeeze(data, axis=np_axis) np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01) verify_squeeze((1, 3, 2, 5), "float32", None) From 36558823974405be67bb83f7824e1c3b29b84e7b Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Mon, 19 Nov 2018 18:34:48 -0800 Subject: [PATCH 6/7] -1 not permitted for dimensions --- tests/python/relay/test_op_level3.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 2862cfe10ac1..890a760f4d86 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -75,7 +75,6 @@ def verify_squeeze(shape, dtype, axis): verify_squeeze((1, 3, 2, 5), "float32", None) verify_squeeze((1, 3, 1), "float32", [0]) - verify_squeeze((1, 3, 2, 5, 1), "float32", [-1]) verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) From aace0d3453ad6ec835359302b3d3b423a5621559 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Tue, 20 Nov 2018 03:24:50 -0500 Subject: [PATCH 7/7] Invoke test_expand_dims() from __main__() --- tests/python/relay/test_op_level1.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index 7e570862d230..d28aa0a56941 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -311,6 +311,7 @@ def test_dense(): test_binary_op() test_expand_dims_infer_type() test_concatenate() + test_expand_dims() test_softmax() test_log_softmax() test_dropout()