From 99cf37619e370777590fb2c8cd991e811ad4ee90 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 15:14:17 +0900 Subject: [PATCH 1/9] add test --- tests/python/frontend/pytorch/test_forward.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 0031f4143fab..2a874bfe7856 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -3992,5 +3992,17 @@ def test_fn(out_int32=False, right=False): verify_model(test_fn(out_int32=True, right=True), [values, boundaries]) +@tvm.testing.uses_gpu +def test_roll(): + def test_fn(shifts, dims): + return lambda x: torch.roll(x, shifts, dims) + + x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]).view(4, 2) + verify_model(test_fn(1, 0), [x]) + verify_model(test_fn(-1, 0), [x]) + verify_model(test_fn(shifts=(2, 1), dims=(0, 1)), [x]) + + if __name__ == "__main__": - pytest.main([__file__]) + # pytest.main([__file__]) + test_roll() From 83b41924ed0c5b8f54dadff278155dc32a3afe99 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 16:13:49 +0900 Subject: [PATCH 2/9] first impl --- python/tvm/relay/frontend/pytorch.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 3fc202a7cc91..965774517d80 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2794,6 +2794,27 @@ def searchsorted(self, inputs, input_types): def bucketize(self, inputs, input_types): return self.searchsorted_common(inputs[1], inputs[0], inputs[2], inputs[3]) + def roll(self, inputs, input_types): + def swap_axes(inp, shape, ax1, ax2): + axes = list(range(len(shape))) + axes[ax1] = ax2 + axes[ax2] = ax1 + return _op.transpose(inp, axes) + + x = inputs[0] + shifts = inputs[1] + dims = inputs[2] + shape = self.infer_shape(x) + start = _expr.const(0, "int64") + roll_dim = _expr.const(shape[dims[0]], "int64") + step = _expr.const(1, "int64") + indices_1d = _op.mod( + _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[0], "int64"), + roll_dim, + ) + indices = swap_axes(_op.tile(indices_1d, shape[:-1] + (1,)), shape, roll_dim, -1) + return _op.gather(x, roll_dim, indices) + # Operator mappings def create_convert_map(self): self.convert_map = { @@ -3021,6 +3042,7 @@ def create_convert_map(self): "aten::any": functools.partial(self.all_any_common, _op.any), "aten::searchsorted": self.searchsorted, "aten::bucketize": self.bucketize, + "aten::roll": self.roll, } def update_convert_map(self, custom_map): From ad9fb0bcf9ef9b62a0d304d832eb77c00aecf818 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 16:26:52 +0900 Subject: [PATCH 3/9] basic example working --- python/tvm/relay/frontend/pytorch.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 965774517d80..b3d16922b623 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2809,11 +2809,13 @@ def swap_axes(inp, shape, ax1, ax2): roll_dim = _expr.const(shape[dims[0]], "int64") step = _expr.const(1, "int64") indices_1d = _op.mod( - _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[0], "int64"), + _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[0], "int64") + roll_dim, roll_dim, ) - indices = swap_axes(_op.tile(indices_1d, shape[:-1] + (1,)), shape, roll_dim, -1) - return _op.gather(x, roll_dim, indices) + indices = swap_axes(_op.tile(indices_1d, shape[:dims[0]] + shape[dims[0]+1:] + (1,)), shape, dims[0], -1) + print("indices shape", self.infer_shape(indices), shape, roll_dim) + print(_infer_value(indices, {})) + return _op.gather(x, dims[0], indices) # Operator mappings def create_convert_map(self): From 85c3ba3d0e46d83205299f3f79bf736f98c7642f Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 16:33:08 +0900 Subject: [PATCH 4/9] all test cases working --- python/tvm/relay/frontend/pytorch.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index b3d16922b623..daae405d3403 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2806,16 +2806,26 @@ def swap_axes(inp, shape, ax1, ax2): dims = inputs[2] shape = self.infer_shape(x) start = _expr.const(0, "int64") - roll_dim = _expr.const(shape[dims[0]], "int64") step = _expr.const(1, "int64") - indices_1d = _op.mod( - _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[0], "int64") + roll_dim, - roll_dim, - ) - indices = swap_axes(_op.tile(indices_1d, shape[:dims[0]] + shape[dims[0]+1:] + (1,)), shape, dims[0], -1) - print("indices shape", self.infer_shape(indices), shape, roll_dim) - print(_infer_value(indices, {})) - return _op.gather(x, dims[0], indices) + + out = x + for i in range(len(dims)): + roll_dim = _expr.const(shape[dims[i]], "int64") + indices_1d = _op.mod( + _op.transform.arange(start, roll_dim, step, "int64") + - _expr.const(shifts[i], "int64") + + roll_dim, + roll_dim, + ) + indices = swap_axes( + _op.tile(indices_1d, shape[: dims[i]] + shape[dims[i] + 1 :] + (1,)), + shape, + dims[i], + -1, + ) + out = _op.gather(out, dims[i], indices) + + return out # Operator mappings def create_convert_map(self): From 210f1185ba0b1a96eb0c2ee9f896d2af314bdcae Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 16:44:30 +0900 Subject: [PATCH 5/9] support adaptive avg and max pool --- python/tvm/relay/frontend/pytorch.py | 37 +++++++++++++++---- tests/python/frontend/pytorch/test_forward.py | 21 ++++++++++- 2 files changed, 48 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index daae405d3403..c919e6a28222 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -849,6 +849,18 @@ def hard_swish(self, inputs, input_types): data = inputs[0] return data * self.hard_sigmoid(inputs, input_types) + def adaptive_avg_pool_1d(self, inputs, input_types): + data = inputs[0] + output_size = inputs[1] + + def func(x): + return _op.nn.adaptive_avg_pool1d(x, output_size=output_size) + + if self.is_quantized_tensor(data): + return qnn_torch.apply_with_upcast(data, func) + + return func(data) + def adaptive_avg_pool_2d(self, inputs, input_types): data = inputs[0] output_size = inputs[1] @@ -861,23 +873,28 @@ def func(x): return func(data) - def adaptive_max_pool_2d(self, inputs, input_types): + def adaptive_avg_pool_3d(self, inputs, input_types): data = inputs[0] output_size = inputs[1] + return _op.nn.adaptive_avg_pool3d(data, output_size=output_size) + def adaptive_max_pool_1d(self, inputs, input_types): + data = inputs[0] + output_size = inputs[1] # returns dummy indices too - return _op.nn.adaptive_max_pool2d(data, output_size=output_size), None + return _op.nn.adaptive_max_pool1d(data, output_size=output_size), None - def adaptive_max_pool_3d(self, inputs, input_types): + def adaptive_max_pool_2d(self, inputs, input_types): data = inputs[0] output_size = inputs[1] # returns dummy indices too - return _op.nn.adaptive_max_pool3d(data, output_size=output_size), None + return _op.nn.adaptive_max_pool2d(data, output_size=output_size), None - def adaptive_avg_pool_3d(self, inputs, input_types): + def adaptive_max_pool_3d(self, inputs, input_types): data = inputs[0] output_size = inputs[1] - return _op.nn.adaptive_avg_pool3d(data, output_size=output_size) + # returns dummy indices too + return _op.nn.adaptive_max_pool3d(data, output_size=output_size), None @staticmethod def convert_const_list(data): @@ -2884,9 +2901,14 @@ def create_convert_map(self): "aten::gelu": self.gelu, "aten::selu": self.selu, "aten::silu": self.silu, + "aten::silu_": self.silu, "aten::log_sigmoid": self.log_sigmoid, + "aten::adaptive_avg_pool1d": self.adaptive_avg_pool_1d, + "aten::adaptive_max_pool1d": self.adaptive_max_pool_1d, "aten::adaptive_avg_pool2d": self.adaptive_avg_pool_2d, "aten::adaptive_max_pool2d": self.adaptive_max_pool_2d, + "aten::adaptive_avg_pool3d": self.adaptive_avg_pool_3d, + "aten::adaptive_max_pool3d": self.adaptive_max_pool_3d, "aten::max_pool2d": self.maxpool_2d, "aten::max_pool2d_with_indices": self.maxpool_2d_with_indices, "aten::max_pool1d": self.maxpool_1d, @@ -2972,6 +2994,7 @@ def create_convert_map(self): "aten::rsqrt": self.make_unary("rsqrt"), "aten::ceil": self.make_unary("ceil"), "aten::floor": self.make_unary("floor"), + "aten::floor_": self.make_unary("floor"), "aten::round": self.make_unary("round"), "aten::isfinite": self.make_unary("isfinite"), "aten::isinf": self.make_unary("isinf"), @@ -2997,8 +3020,6 @@ def create_convert_map(self): "aten::bitwise_xor": self.bitwise_xor, "aten::Bool": self.Bool, "aten::Float": self.Float, - "aten::adaptive_avg_pool3d": self.adaptive_avg_pool_3d, - "aten::adaptive_max_pool3d": self.adaptive_max_pool_3d, "aten::rsub": self.rsub, "aten::embedding": self.embedding, "aten::one_hot": self.one_hot, diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 2a874bfe7856..5299125c63cd 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -735,13 +735,30 @@ def test_forward_log_sigmoid(): @tvm.testing.uses_gpu -def test_forward_adaptiveavgpool(): +def test_forward_adaptive_avgpool(): torch.set_grad_enabled(False) input_shape = [1, 3, 10, 10] input_data = torch.rand(input_shape).float() verify_model(torch.nn.AdaptiveAvgPool2d([1, 1]).eval(), input_data=input_data) verify_model(torch.nn.AdaptiveAvgPool2d([10, 10]).eval(), input_data=input_data) + input_data = torch.rand([1, 3, 10]).float() + verify_model(torch.nn.AdaptiveAvgPool1d([1]).eval(), input_data=input_data) + verify_model(torch.nn.AdaptiveAvgPool1d([5]).eval(), input_data=input_data) + + +@tvm.testing.uses_gpu +def test_forward_adaptive_maxpool(): + torch.set_grad_enabled(False) + input_shape = [1, 3, 10, 10] + input_data = torch.rand(input_shape).float() + verify_model(torch.nn.AdaptiveMaxPool2d([1, 1]).eval(), input_data=input_data) + verify_model(torch.nn.AdaptiveMaxPool2d([10, 10]).eval(), input_data=input_data) + + input_data = torch.rand([1, 3, 10]).float() + verify_model(torch.nn.AdaptiveMaxPool1d([1]).eval(), input_data=input_data) + verify_model(torch.nn.AdaptiveMaxPool1d([5]).eval(), input_data=input_data) + @tvm.testing.uses_gpu def test_forward_maxpool2d(): @@ -4005,4 +4022,4 @@ def test_fn(shifts, dims): if __name__ == "__main__": # pytest.main([__file__]) - test_roll() + test_forward_adaptive_maxpool() From 8f32bb02536411d65cc2b08882b56dee3103e39b Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 16:52:28 +0900 Subject: [PATCH 6/9] cleanup --- python/tvm/relay/frontend/pytorch.py | 61 +++++++------------ tests/python/frontend/pytorch/test_forward.py | 3 +- 2 files changed, 23 insertions(+), 41 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index c919e6a28222..2dc625a36792 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -849,52 +849,23 @@ def hard_swish(self, inputs, input_types): data = inputs[0] return data * self.hard_sigmoid(inputs, input_types) - def adaptive_avg_pool_1d(self, inputs, input_types): + def adaptive_avg_pool(self, op, inputs, input_types): data = inputs[0] output_size = inputs[1] def func(x): - return _op.nn.adaptive_avg_pool1d(x, output_size=output_size) + return op(x, output_size=output_size) if self.is_quantized_tensor(data): return qnn_torch.apply_with_upcast(data, func) return func(data) - def adaptive_avg_pool_2d(self, inputs, input_types): - data = inputs[0] - output_size = inputs[1] - - def func(x): - return _op.nn.adaptive_avg_pool2d(x, output_size=output_size) - - if self.is_quantized_tensor(data): - return qnn_torch.apply_with_upcast(data, func) - - return func(data) - - def adaptive_avg_pool_3d(self, inputs, input_types): - data = inputs[0] - output_size = inputs[1] - return _op.nn.adaptive_avg_pool3d(data, output_size=output_size) - - def adaptive_max_pool_1d(self, inputs, input_types): - data = inputs[0] - output_size = inputs[1] - # returns dummy indices too - return _op.nn.adaptive_max_pool1d(data, output_size=output_size), None - - def adaptive_max_pool_2d(self, inputs, input_types): - data = inputs[0] - output_size = inputs[1] - # returns dummy indices too - return _op.nn.adaptive_max_pool2d(data, output_size=output_size), None - - def adaptive_max_pool_3d(self, inputs, input_types): + def adaptive_max_pool(self, op, inputs, input_types): data = inputs[0] output_size = inputs[1] # returns dummy indices too - return _op.nn.adaptive_max_pool3d(data, output_size=output_size), None + return op(data, output_size=output_size), None @staticmethod def convert_const_list(data): @@ -2903,12 +2874,24 @@ def create_convert_map(self): "aten::silu": self.silu, "aten::silu_": self.silu, "aten::log_sigmoid": self.log_sigmoid, - "aten::adaptive_avg_pool1d": self.adaptive_avg_pool_1d, - "aten::adaptive_max_pool1d": self.adaptive_max_pool_1d, - "aten::adaptive_avg_pool2d": self.adaptive_avg_pool_2d, - "aten::adaptive_max_pool2d": self.adaptive_max_pool_2d, - "aten::adaptive_avg_pool3d": self.adaptive_avg_pool_3d, - "aten::adaptive_max_pool3d": self.adaptive_max_pool_3d, + "aten::adaptive_avg_pool1d": functools.partial( + self.adaptive_avg_pool, _op.nn.adaptive_avg_pool1d + ), + "aten::adaptive_avg_pool2d": functools.partial( + self.adaptive_avg_pool, _op.nn.adaptive_avg_pool2d + ), + "aten::adaptive_avg_pool3d": functools.partial( + self.adaptive_avg_pool, _op.nn.adaptive_avg_pool3d + ), + "aten::adaptive_max_pool1d": functools.partial( + self.adaptive_max_pool, _op.nn.adaptive_max_pool1d + ), + "aten::adaptive_max_pool2d": functools.partial( + self.adaptive_max_pool, _op.nn.adaptive_max_pool2d + ), + "aten::adaptive_max_pool3d": functools.partial( + self.adaptive_max_pool, _op.nn.adaptive_max_pool3d + ), "aten::max_pool2d": self.maxpool_2d, "aten::max_pool2d_with_indices": self.maxpool_2d_with_indices, "aten::max_pool1d": self.maxpool_1d, diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 5299125c63cd..5057f0d2b6b8 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -4021,5 +4021,4 @@ def test_fn(shifts, dims): if __name__ == "__main__": - # pytest.main([__file__]) - test_forward_adaptive_maxpool() + pytest.main([__file__]) From 30ea0fc970cce14b32a923840a2de8309eef505f Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 17:25:45 +0900 Subject: [PATCH 7/9] axes transpose logic fixed for roll --- python/tvm/relay/frontend/pytorch.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 2dc625a36792..fb5d219996f1 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2783,10 +2783,9 @@ def bucketize(self, inputs, input_types): return self.searchsorted_common(inputs[1], inputs[0], inputs[2], inputs[3]) def roll(self, inputs, input_types): - def swap_axes(inp, shape, ax1, ax2): + def slide_axes(inp, shape, ax): axes = list(range(len(shape))) - axes[ax1] = ax2 - axes[ax2] = ax1 + axes = axes[:ax] + [-1] + axes[ax:-1] return _op.transpose(inp, axes) x = inputs[0] @@ -2805,11 +2804,12 @@ def swap_axes(inp, shape, ax1, ax2): + roll_dim, roll_dim, ) - indices = swap_axes( + # First fill in the last axis with roll indices, and then do transpose to + # bring the roll indices into the desired axis. + indices = slide_axes( _op.tile(indices_1d, shape[: dims[i]] + shape[dims[i] + 1 :] + (1,)), shape, dims[i], - -1, ) out = _op.gather(out, dims[i], indices) From 58aaebc4c90aa9829b309f63473a035632616c88 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 17:38:43 +0900 Subject: [PATCH 8/9] pylint --- python/tvm/relay/frontend/pytorch.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index fb5d219996f1..58af8b24b6c0 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2796,8 +2796,8 @@ def slide_axes(inp, shape, ax): step = _expr.const(1, "int64") out = x - for i in range(len(dims)): - roll_dim = _expr.const(shape[dims[i]], "int64") + for i, dim in enumerate(dims): + roll_dim = _expr.const(shape[i], "int64") indices_1d = _op.mod( _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[i], "int64") @@ -2807,11 +2807,11 @@ def slide_axes(inp, shape, ax): # First fill in the last axis with roll indices, and then do transpose to # bring the roll indices into the desired axis. indices = slide_axes( - _op.tile(indices_1d, shape[: dims[i]] + shape[dims[i] + 1 :] + (1,)), + _op.tile(indices_1d, shape[:dim] + shape[dim + 1 :] + (1,)), shape, - dims[i], + dim, ) - out = _op.gather(out, dims[i], indices) + out = _op.gather(out, dim, indices) return out From a6c509b5e9446ca04d58f037f30f5f406b6bfdc3 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 26 Oct 2021 17:43:02 +0900 Subject: [PATCH 9/9] fixed roll dim indexing --- python/tvm/relay/frontend/pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 58af8b24b6c0..13704ff7aad9 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2797,7 +2797,7 @@ def slide_axes(inp, shape, ax): out = x for i, dim in enumerate(dims): - roll_dim = _expr.const(shape[i], "int64") + roll_dim = _expr.const(shape[dim], "int64") indices_1d = _op.mod( _op.transform.arange(start, roll_dim, step, "int64") - _expr.const(shifts[i], "int64")