[fx2trt] Support conv1d converter for module (#45)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/45

Extend the support for conv1d converter for the module for fx2trt
As 2d convolution is not supported by TensorRT, this patch maps conv1d operation to conv2d with a pair of unsqueeze and squeeze.

Reviewed By: wushirong

Differential Revision: D35395174

fbshipit-source-id: 5b564609998f103d083d1451e5d64f6224eaf695

Author: kparichay

fx/converters/convolution.py

@@ -1,4 +1,6 @@
 # @manual=//deeplearning/trt/python:py_tensorrt
+import numpy as np
+
 import tensorrt as trt
 import torch
 from fx2trt_oss.fx.converter_registry import tensorrt_converter
@@ -22,6 +24,20 @@ def common_conv(network, mod, dimension, input_val, layer_name, is_quantized):
     kernel = to_numpy(mod.weight() if is_quantized else mod.weight)
     bias = to_numpy(mod.bias() if is_quantized else mod.bias)
 
+    if dimension == 1:
+        # Append unsqueeze before conv2d to calculate conv1d
+        unsqueeze_layer = network.add_shuffle(input=input_val)
+        unsqueeze_layer.reshape_dims = (*input_val.shape, 1)
+        unsqueeze_layer.name = f"{layer_name}_unsqueeze"
+        input_val = unsqueeze_layer.get_output(0)
+
+        padding = padding + (0,)
+        kernel = np.expand_dims(kernel, -1)
+        kernel_size = kernel.shape[2:]
+        if bias is not None:
+            bias = bias[None]
+            # bias = np.expand_dims(bias, -1)
+
     layer = network.add_convolution_nd(
         input=input_val,
         num_output_maps=mod.out_channels,
@@ -39,7 +55,15 @@ def common_conv(network, mod, dimension, input_val, layer_name, is_quantized):
         # Assume the dtype of activation is torch.quint8
         mark_as_int8_layer(layer, get_dyn_range(mod.scale, mod.zero_point, torch.quint8))
 
-    return layer.get_output(0)
+    result = layer.get_output(0)
+    if dimension == 1:
+        # Append squeeze after conv2d to calculate conv1d
+        squeeze_layer = network.add_shuffle(input=result)
+        squeeze_layer.reshape_dims = tuple(result.shape[:-1])
+        squeeze_layer.name = f"{layer_name}_squeeze"
+        result = squeeze_layer.get_output(0)
+
+    return result
 
 
 def common_conv_relu(network, mod, dimension, input_val, layer_name, is_quantized):
@@ -62,6 +86,20 @@ def common_conv_relu(network, mod, dimension, input_val, layer_name, is_quantize
     return layer.get_output(0)
 
 
+@tensorrt_converter(torch.nn.modules.conv.Conv1d)
+def conv1d(network, submod, args, kwargs, layer_name):
+    # args/kwargs should have already been normalized to kwargs
+    assert len(args) == 0
+    input_val = kwargs["input"]
+
+    if not isinstance(input_val, trt.tensorrt.ITensor):
+        raise RuntimeError(f"Conv1d received input {input_val} that is not part "
+                           "of the TensorRT region!")
+
+    if layer_name is None:
+        raise RuntimeError("layer name is none")
+    return common_conv(network, submod, dimension=1, input_val=input_val, layer_name=layer_name, is_quantized=False)
+
 @tensorrt_converter(torch.nn.modules.conv.Conv2d)
 def conv2d(network, submod, args, kwargs, layer_name):
     # args/kwargs should have already been normalized to kwargs

test/converters/vanilla/test_convolution.py

@@ -8,6 +8,39 @@
 
 
 class TestConvolutionConverter(VanillaTestCase):
+    @parameterized.expand(
+        [
+            ("default", 1),
+            param("no_bias", 1, bias=False),
+            ("tuple_parameters", 1, (1), (0)),
+            param("non_zero_padding", 1, padding=1),
+            param("dilation", 1, dilation=2),
+            param("groups", 1, groups=3),
+        ]
+    )
+    def test_conv1d(
+        self,
+        test_name,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv1d(
+                    3, 6, kernel_size, stride, padding, dilation, groups, bias
+                )
+
+            def forward(self, x):
+                return self.conv(x)
+
+        inputs = [torch.randn(1, 3, 224)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.nn.modules.conv.Conv1d})
+
     @parameterized.expand(
         [
             ("default", 1),