From 86588e5ba3af1635ede5039de3c79a5e0875d0ec Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Thu, 18 Mar 2021 11:59:41 +0100 Subject: [PATCH 1/4] [TORCH] Implement avg_pool1d --- python/tvm/relay/frontend/pytorch.py | 25 +++++++++++++++++++ tests/python/frontend/pytorch/test_forward.py | 19 ++++++++++++-- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 8ae1e862ffd5..4a80fcd00ee4 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -1353,6 +1353,30 @@ def softplus(self, inputs, input_types): beta = _expr.const(float(inputs[1]), dtype=dtype) return _op.log(_op.exp(inputs[0] * beta) + _expr.const(1.0, dtype=dtype)) / beta + def avg_pool1d(self, inputs, input_types): + data = inputs[0] + + pool_size = self.convert_const_list(inputs[1]) + strides = self.convert_const_list(inputs[2] if inputs[2] else pool_size) + padding = inputs[3] + ceil_mode = int(inputs[4]) + count_include_pad = int(inputs[5]) + + def func(x): + return _op.nn.avg_pool1d( + x, + pool_size=pool_size, + strides=strides, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + + if self.is_quantized_tensor(data): + return qnn_torch.apply_with_upcast(data, func) + + return func(data) + def avg_pool2d(self, inputs, input_types): data = inputs[0] @@ -2350,6 +2374,7 @@ def create_convert_map(self): "aten::log_softmax": self.log_softmax, "aten::sigmoid": self.sigmoid, "aten::softplus": self.softplus, + "aten::avg_pool1d": self.avg_pool1d, "aten::avg_pool2d": self.avg_pool2d, "aten::avg_pool3d": self.avg_pool3d, "aten::linear": self.linear, diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index d0edfd9c8036..ad9d15fc53fb 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -809,7 +809,21 @@ def forward(self, *args): @tvm.testing.uses_gpu -def test_forward_avgpool(): +def test_forward_avgpool1d(): + torch.set_grad_enabled(False) + input_shape = [1, 3, 10] + + class AvgPool1D2(Module): + def forward(self, *args): + return torch.nn.functional.avg_pool1d(args[0], kernel_size=[10]) + + input_data = torch.rand(input_shape).float() + verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data) + verify_model(AvgPool1D2().float().eval(), input_data=input_data) + + +@tvm.testing.uses_gpu +def test_forward_avgpool2d(): torch.set_grad_enabled(False) input_shape = [1, 3, 10, 10] @@ -3838,7 +3852,8 @@ def test_fn(is_sorted, return_inverse, return_counts): test_forward_logsoftmax() test_forward_sigmoid() test_forward_dense() - test_forward_avgpool() + test_forward_avgpool1d() + test_forward_avgpool2d() test_forward_avgpool3d() test_forward_dropout() test_forward_slice() From ba885fc348248d05273523b39d533d8ab28425b6 Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Thu, 18 Mar 2021 21:56:33 +0100 Subject: [PATCH 2/4] [TORCH] Unify creation of avg_pooling operations --- python/tvm/relay/frontend/pytorch.py | 110 ++++++++++++--------------- 1 file changed, 47 insertions(+), 63 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 4a80fcd00ee4..b64afb70a2ee 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -1353,71 +1353,55 @@ def softplus(self, inputs, input_types): beta = _expr.const(float(inputs[1]), dtype=dtype) return _op.log(_op.exp(inputs[0] * beta) + _expr.const(1.0, dtype=dtype)) / beta - def avg_pool1d(self, inputs, input_types): - data = inputs[0] - - pool_size = self.convert_const_list(inputs[1]) - strides = self.convert_const_list(inputs[2] if inputs[2] else pool_size) - padding = inputs[3] - ceil_mode = int(inputs[4]) - count_include_pad = int(inputs[5]) - - def func(x): - return _op.nn.avg_pool1d( - x, - pool_size=pool_size, - strides=strides, - padding=padding, - ceil_mode=ceil_mode, - count_include_pad=count_include_pad, - ) - - if self.is_quantized_tensor(data): - return qnn_torch.apply_with_upcast(data, func) - - return func(data) - - def avg_pool2d(self, inputs, input_types): - data = inputs[0] - - pool_size = self.convert_const_list(inputs[1]) - strides = self.convert_const_list(inputs[2] if inputs[2] else pool_size) - padding = inputs[3] - ceil_mode = int(inputs[4]) - count_include_pad = int(inputs[5]) - - def func(x): - return _op.nn.avg_pool2d( - x, - pool_size=pool_size, - strides=strides, - padding=padding, - ceil_mode=ceil_mode, - count_include_pad=count_include_pad, - ) + def make_avg_pool(self, dim): + + def avg_pool(inputs, input_types): + data = inputs[0] - if self.is_quantized_tensor(data): - return qnn_torch.apply_with_upcast(data, func) + pool_size = self.convert_const_list(inputs[1]) + strides = self.convert_const_list(inputs[2] if inputs[2] else pool_size) + padding = inputs[3] + ceil_mode = int(inputs[4]) + count_include_pad = int(inputs[5]) - return func(data) + def func(x): + if dim == 1: + return _op.nn.avg_pool1d( + x, + pool_size=pool_size, + strides=strides, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + elif dim == 2: + return _op.nn.avg_pool2d( + x, + pool_size=pool_size, + strides=strides, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + elif dim == 3: + return _op.nn.avg_pool3d( + x, + pool_size=pool_size, + strides=strides, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + else: + msg = "Average Pooling dimension should be between 1 and 3" + raise RuntimeError(msg) - def avg_pool3d(self, inputs, input_types): - data = inputs[0] + if self.is_quantized_tensor(data): + return qnn_torch.apply_with_upcast(data, func) - pool_size = inputs[1] - strides = inputs[2] if inputs[2] else pool_size - padding = inputs[3] - ceil_mode = int(inputs[4]) - count_include_pad = int(inputs[5]) + return func(data) - return _op.nn.avg_pool3d( - data, - pool_size=pool_size, - strides=strides, - padding=padding, - ceil_mode=ceil_mode, - count_include_pad=count_include_pad, - ) + return avg_pool def linear(self, inputs, input_types): # https://pytorch.org/docs/stable/nn.functional.html#linear @@ -2374,9 +2358,9 @@ def create_convert_map(self): "aten::log_softmax": self.log_softmax, "aten::sigmoid": self.sigmoid, "aten::softplus": self.softplus, - "aten::avg_pool1d": self.avg_pool1d, - "aten::avg_pool2d": self.avg_pool2d, - "aten::avg_pool3d": self.avg_pool3d, + "aten::avg_pool1d": self.make_avg_pool(1), + "aten::avg_pool2d": self.make_avg_pool(2), + "aten::avg_pool3d": self.make_avg_pool(3), "aten::linear": self.linear, "aten::dropout": self.dropout, "aten::dropout_": self.dropout, From 1835aa6289d3e1ef3e3c00db0b92e3d3424132bb Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Thu, 18 Mar 2021 22:09:33 +0100 Subject: [PATCH 3/4] [TORCH] Add tests for avg pooling with padding --- tests/python/frontend/pytorch/test_forward.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index ad9d15fc53fb..3482965cfb58 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -820,6 +820,7 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data) verify_model(AvgPool1D2().float().eval(), input_data=input_data) + verify_model(torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data) @tvm.testing.uses_gpu @@ -834,6 +835,7 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool2d(kernel_size=[10, 10]).eval(), input_data=input_data) verify_model(AvgPool2D2().float().eval(), input_data=input_data) + verify_model(torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data) @tvm.testing.uses_gpu @@ -848,6 +850,7 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool3d(kernel_size=[10, 10, 10]).eval(), input_data=input_data) verify_model(AvgPool3D1().float().eval(), input_data=input_data) + verify_model(torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data) @tvm.testing.uses_gpu From 2c2a843c555a116d3f466cfcd1353219f99c643e Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Thu, 18 Mar 2021 22:29:44 +0100 Subject: [PATCH 4/4] [TORCH] Make format checks happy with unified avg_pool --- python/tvm/relay/frontend/pytorch.py | 1 - tests/python/frontend/pytorch/test_forward.py | 12 +++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index b64afb70a2ee..cb9ea6a043f4 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -1354,7 +1354,6 @@ def softplus(self, inputs, input_types): return _op.log(_op.exp(inputs[0] * beta) + _expr.const(1.0, dtype=dtype)) / beta def make_avg_pool(self, dim): - def avg_pool(inputs, input_types): data = inputs[0] diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 3482965cfb58..572aa472c540 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -820,7 +820,9 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data) verify_model(AvgPool1D2().float().eval(), input_data=input_data) - verify_model(torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data) + verify_model( + torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data + ) @tvm.testing.uses_gpu @@ -835,7 +837,9 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool2d(kernel_size=[10, 10]).eval(), input_data=input_data) verify_model(AvgPool2D2().float().eval(), input_data=input_data) - verify_model(torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data) + verify_model( + torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data + ) @tvm.testing.uses_gpu @@ -850,7 +854,9 @@ def forward(self, *args): input_data = torch.rand(input_shape).float() verify_model(torch.nn.AvgPool3d(kernel_size=[10, 10, 10]).eval(), input_data=input_data) verify_model(AvgPool3D1().float().eval(), input_data=input_data) - verify_model(torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data) + verify_model( + torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data + ) @tvm.testing.uses_gpu