[fx2trt] avg_pool, max_pool ops improvement

Summary:
nn.avg_pool1d,nn.avg_pool2d,nn.max_pool1d,nn.max_pool2d,
will set stride=kernel_size if stride=None
But nn.functional.avg_pool and nn.functional.max_pool does not set that. This will set stride=Node to TRT and leads to a TRT error.
1. This diff fix this hole.
2. Add more test cases
3. Fix LearningToPaint model in torchbench

Reviewed By: yinghai

Differential Revision: D34953195

fbshipit-source-id: 0eb73de29bc7461be1c7c6fa96f820829b2cce10

Author: Wei Wei

fx/converters/acc_ops_converters.py

@@ -1121,24 +1121,28 @@ def acc_ops_max_pool1d(
 
     input_trt = unsqueeze_layer.get_output(0)
 
-    kernel_size = kwargs["kernel_size"]
-    stride = kwargs["stride"]
-    padding = kwargs["padding"]
-    dilation = kwargs["dilation"]
+    kernel_size = extend_attr_to_tuple(kwargs["kernel_size"], 1)
+    stride = extend_attr_to_tuple(kwargs["stride"], 1)
+    padding = extend_attr_to_tuple(kwargs["padding"], 1)
+    dilation = extend_attr_to_tuple(kwargs["dilation"], 1)
+
     ceil_mode = kwargs["ceil_mode"]
 
-    if any([not isinstance(param, int) for param in  [kernel_size, stride, padding, dilation]]):
+    if len(stride) == 0 or stride[0] == None:
+        stride = kernel_size
+
+    if any([not isinstance(param, int) for param in  [kernel_size[0], stride[0], padding[0], dilation[0]]]):
         raise RuntimeError(f"Parameters kernel_size, stride, padding, and dilation should be of type int.")
-    if dilation != 1:
+    if dilation[0] != 1:
         raise RuntimeError(
             f"Only support dilation=1 for maxpool, but got {dilation}"
         )
 
     max_pooling_layer = network.add_pooling(
-        input=input_trt, type=trt.PoolingType.MAX, window_size=(kernel_size, 1)
+        input=input_trt, type=trt.PoolingType.MAX, window_size=(kernel_size[0], 1)
     )
-    max_pooling_layer.stride_nd = (stride, 1)
-    max_pooling_layer.padding_nd =  (padding, 0)
+    max_pooling_layer.stride_nd = stride +(1,)
+    max_pooling_layer.padding_nd = padding + (0,)
     set_layer_name(max_pooling_layer, target, name)
 
     if ceil_mode:
@@ -1171,6 +1175,9 @@ def acc_ops_max_pool2d(
     dilation = extend_attr_to_tuple(kwargs["dilation"], 2)
     ceil_mode = kwargs["ceil_mode"]
 
+    if len(stride) == 0 or stride[0] == None:
+        stride = kernel_size
+
     if dilation != (1, 1):
         raise RuntimeError(
             f"Only support dilation=(1, 1) for maxpool, but got {dilation}"
@@ -1445,6 +1452,9 @@ def acc_ops_avg_pool1d(
     ceil_mode = kwargs["ceil_mode"]
     count_include_pad = kwargs["count_include_pad"]
 
+    if len(stride) == 0 or stride[0] == None:
+        stride = kernel_size
+
     shuffle_layer = network.add_shuffle(input_val)
     shuffle_layer.reshape_dims = tuple(input_val.shape) + (1,)
     set_layer_name(shuffle_layer, target, name + "_shuffle1")
@@ -1493,6 +1503,9 @@ def acc_ops_avg_pool2d(
     count_include_pad = kwargs["count_include_pad"]
     divisor_override = kwargs["divisor_override"]
 
+    if len(stride) == 0 or stride[0] == None:
+        stride = kernel_size
+
     if divisor_override:
         raise RuntimeError("TensorRT does not support divisor_override.")
 

test/converters/acc_op/test_avgpool.py

@@ -84,5 +84,70 @@ def forward(self, x):
         inputs = [torch.randn(1, 3, 224, 224)]
         self.run_test(TestModule(), inputs, expected_ops={acc_ops.avg_pool2d})
 
+    @parameterized.expand(
+        [
+            ("kernal_size", 1),
+            param("stride", 2, stride=()),
+        ]
+    )
+    def test_stride_none__avg_pool1d(
+        self,
+        test_name,
+        kernel_size,
+        stride=None,
+        padding=0,
+        ceil_mode=False,
+        count_include_pad=True
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.nn.functional.avg_pool1d(
+                    x,
+                    kernel_size,
+                    stride=stride,
+                    padding=padding,
+                    ceil_mode=ceil_mode,
+                    count_include_pad=count_include_pad,
+                )
+
+        inputs = [torch.randn(1, 3, 224)]
+        self.run_test(TestModule(), inputs, expected_ops={acc_ops.avg_pool1d})
+
+    @parameterized.expand(
+        [
+            ("kernal_size", 2),
+            param("stride", 2, stride=()),
+        ]
+    )
+    def test_stride_none_avg_pool2d(
+        self,
+        test_name,
+        kernel_size,
+        stride=None,
+        padding=0,
+        ceil_mode=False,
+        count_include_pad=True,
+        divisor_override=None,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+            def forward(self, x):
+                return torch.nn.functional.avg_pool2d(
+                    x,
+                    kernel_size,
+                    stride=stride,
+                    padding=padding,
+                    ceil_mode=ceil_mode,
+                    count_include_pad=count_include_pad,
+                    divisor_override=divisor_override,
+                )
+
+        inputs = [torch.randn(1, 3, 224, 224)]
+        self.run_test(TestModule(), inputs, expected_ops={acc_ops.avg_pool2d})
+
 if __name__ == '__main__':
     run_tests()

test/converters/acc_op/test_maxpool.py

@@ -92,5 +92,58 @@ def forward(self, x):
             TestModule(), input_specs, expected_ops={acc_ops.max_pool2d}
         )
 
+    @parameterized.expand(
+        [
+            ("default", 1),
+            param("stride", 2, stride=()),
+        ]
+    )
+    def test_stride_none_max_pool1d(self,
+        test_name,
+        kernel_size,
+        stride=None,
+        padding=0,
+        dilation=1,
+        ceil_mode=False,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.nn.functional.max_pool1d(
+                    x, kernel_size, stride=stride, padding=padding, ceil_mode=ceil_mode, dilation=dilation
+                )
+
+        inputs = [torch.randn(1, 3, 224)]
+        self.run_test(TestModule(), inputs, expected_ops={acc_ops.max_pool1d}, test_explicit_batch_dim=False,)
+
+
+    @parameterized.expand(
+        [
+            ("default", 1),
+            param("stride", 2, stride=()),
+        ]
+    )
+    def test_stride_none_max_pool2d(
+        self,
+        test_name,
+        kernel_size,
+        stride=None,
+        padding=0,
+        ceil_mode=False,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return  torch.nn.functional.max_pool2d(
+                    x, kernel_size, stride=stride, padding=padding, ceil_mode=ceil_mode
+                )
+
+        inputs = [torch.randn(1, 3, 224, 224)]
+        self.run_test(TestModule(), inputs, expected_ops={acc_ops.max_pool2d})
+
 if __name__ == '__main__':
     run_tests()