Add support for torch.nn.functional.max_pool1d in fx2trt. (#15)

Summary:
Pull Request resolved: pytorch/fx2trt#15

This commit adds support of max_pool1d in torch fx2trt

Reviewed By: 842974287

Differential Revision: D34695867

fbshipit-source-id: 5f06e526edf4cdd5fe8ffee022802aade5808715

Author: EgorLakomkin

fx/converters/acc_ops_converters.py

@@ -1101,6 +1101,53 @@ def acc_ops_fmod(
     )
     return sub_value
 
+@tensorrt_converter(acc_ops.max_pool1d, no_explicit_batch_dim=True)
+def acc_ops_max_pool1d(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    if not network.has_implicit_batch_dimension:
+        raise RuntimeError("Current implementation does not support dynamic shape. Make sure that the network has an explicit batch dimension!")
+
+    input_trt = kwargs["input"]
+
+    # adds unsqueeze layer -> max pool 2d -> squeeze layer to emulate max pool 1d.
+    unsqueeze_layer = network.add_shuffle(input=input_trt)
+    unsqueeze_layer.reshape_dims = tuple([*input_trt.shape, 1])
+    set_layer_name(unsqueeze_layer, target, name+"_unsqueeze")
+
+    input_trt = unsqueeze_layer.get_output(0)
+
+    kernel_size = kwargs["kernel_size"]
+    stride = kwargs["stride"]
+    padding = kwargs["padding"]
+    dilation = kwargs["dilation"]
+    ceil_mode = kwargs["ceil_mode"]
+
+    if any([not isinstance(param, int) for param in  [kernel_size, stride, padding, dilation]]):
+        raise RuntimeError(f"Parameters kernel_size, stride, padding, and dilation should be of type int.")
+    if dilation != 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)
+    )
+    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:
+        max_pooling_layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
+    input_trt = max_pooling_layer.get_output(0)
+    squeeze_layer = network.add_shuffle(input=input_trt)
+    squeeze_layer.reshape_dims = tuple(input_trt.shape[:-1])
+    set_layer_name(squeeze_layer, target, name+"_squeeze")
+    return squeeze_layer.get_output(0)
 
 @tensorrt_converter(acc_ops.max_pool2d)
 def acc_ops_max_pool2d(

test/converters/acc_op/test_maxpool.py

@@ -8,6 +8,38 @@
 
 
 class TestMaxPoolConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("default", 1),
+            ("kernel_3", 3),
+            ("stride", 1, 2),
+            param("padding", 2, padding=1),
+            param("padding_even", 5, padding=2),
+            param("ceil_mode", 1, ceil_mode=True),
+        ]
+    )
+    def test_max_pool1d(self,
+        test_name,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        ceil_mode=False,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.max_pool = torch.nn.MaxPool1d(
+                    kernel_size, stride, padding, ceil_mode=ceil_mode, dilation=dilation
+                )
+
+            def forward(self, x):
+                return self.max_pool(x)
+
+        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),

test/tracer/test_acc_tracer.py

@@ -2062,6 +2062,7 @@ def test_all_acc_ops_registered(self):
             acc_normalizer._acc_ops,
             {
                 acc_ops.linear,
+                acc_ops.max_pool1d,
                 acc_ops.max_pool2d,
                 acc_ops.flatten,
                 acc_ops.adaptive_avg_pool2d,

tracer/acc_tracer/acc_ops.py

@@ -80,6 +80,21 @@ def squeeze(*, input, dim=None):
     return input.squeeze(dim=dim)
 
 
+@register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool1d))
+@register_acc_op
+def max_pool1d(
+    *, input, kernel_size, stride, padding, dilation, ceil_mode, return_indices
+):
+    return nn.functional.max_pool1d(
+        input=input,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=padding,
+        dilation=dilation,
+        ceil_mode=ceil_mode,
+        return_indices=return_indices,
+    )
+
 @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool2d))
 @register_acc_op
 def max_pool2d(