Add support for torch.nn.functional.conv_transpose3d in fx2trt (#40)

Summary:
Pull Request resolved: pytorch/fx2trt#40

add support for torch.nn.functional.conv_transpose3d

Reviewed By: 842974287

Differential Revision: D35176171

fbshipit-source-id: 5b10fba4568ad9fb26203f3dd396e13a98bc3495

Author: SherlockNoMad

fx/converters/acc_ops_converters.py

@@ -172,7 +172,8 @@ def acc_ops_convnd(
 
 
 @tensorrt_converter(acc_ops.conv_transpose2d)
-def acc_ops_conv_transpose2d(
+@tensorrt_converter(acc_ops.conv_transpose3d)
+def acc_ops_conv_transposend(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
@@ -195,7 +196,7 @@ def acc_ops_conv_transpose2d(
     # right now
     if kwargs["bias"] is not None and not isinstance(kwargs["bias"], torch.Tensor):
         raise RuntimeError(
-            f"conv {name} has bias of type {type(kwargs['bias'])}, Expect Optional[Tensor]"
+            f"ConvTranspose {name} has bias of type {type(kwargs['bias'])}, Expect Optional[Tensor]"
         )
     bias = to_numpy(kwargs["bias"])  # type: ignore[arg-type]
 
@@ -205,10 +206,11 @@ def acc_ops_conv_transpose2d(
         # will need to use uninitialized weight and set it later to support
         # ITensor weights
         dummy_weight = trt.Weights()
-        # nn.ConvTranspose2d weight size is (in_channels, out_channels/groups, kernel_0, kernel_1)
+
+        # nn.ConvTranspose2d/3d weight size is (in_channels, out_channels/groups, kernel_0, kernel_1, [kernel_2])
         layer = network.add_deconvolution_nd(
             input=input_val,
-            num_output_maps=weight.shape[1]*kwargs["groups"],
+            num_output_maps=weight.shape[1] * kwargs["groups"],
             kernel_shape=weight.shape[2:],
             kernel=dummy_weight,
             bias=bias,
@@ -221,10 +223,10 @@ def acc_ops_conv_transpose2d(
                 f"conv {name} has weight of type {type(kwargs['weight'])}, Expect Optional[Tensor]"
             )
         weight = to_numpy(kwargs["weight"])
-        # nn.ConvTranspose2d weight size is (in_channels, out_channels/groups, kernel_0, kernel_1)
+        # nn.ConvTranspose2d/3d weight size is (in_channels, out_channels/groups, kernel_0, kernel_1, [kernel_2])
         layer = network.add_deconvolution_nd(
             input=input_val,
-            num_output_maps=weight.shape[1]*kwargs["groups"],
+            num_output_maps=weight.shape[1] * kwargs["groups"],
             kernel_shape=weight.shape[2:],
             kernel=weight,
             bias=bias,

test/converters/acc_op/test_transpose_convolution.py

@@ -32,7 +32,7 @@ def test_conv_transpose2d(
         class TestModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.conv = torch.nn.ConvTranspose2d(
+                self.conv_transpose = torch.nn.ConvTranspose2d(
                     in_channels=3,
                     out_channels=6,
                     kernel_size=kernel_size,
@@ -45,7 +45,7 @@ def __init__(self):
                 )
 
             def forward(self, x):
-                return self.conv(x)
+                return self.conv_transpose(x)
 
         inputs = [torch.randn(1, 3, 224, 224)]
         self.run_test(TestModule(), inputs, expected_ops={acc_ops.conv_transpose2d})
@@ -54,10 +54,10 @@ def test_conv_transpose2d_with_dynamic_shape(self):
         class TestModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.conv = torch.nn.ConvTranspose2d(3, 3, 1)
+                self.conv_transpose = torch.nn.ConvTranspose2d(3, 3, 1)
 
             def forward(self, x):
-                return self.conv(x)
+                return self.conv_transpose(x)
 
         input_specs = [
             InputTensorSpec(
@@ -71,5 +71,67 @@ def forward(self, x):
         )
 
 
+    @parameterized.expand(
+        [
+            ("default", 1),
+            param("no_bias", 1, bias=False),
+            ("tuple_parameters", 1, (1, 1, 1), (1, 1, 1)),
+            param("non_zero_padding", 1, padding=1),
+            param("dilation", 1, dilation=2),
+            param("groups", 1, groups=3),
+        ]
+    )
+    def test_conv_transpose3d(
+        self,
+        _,
+        kernel_size,
+        stride=1,
+        padding=0,
+        output_padding=0,
+        groups=1,
+        bias=True,
+        dilation=1,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv_transpose = torch.nn.ConvTranspose3d(
+                    in_channels=3,
+                    out_channels=6,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    padding=padding,
+                    output_padding=output_padding,
+                    groups=groups,
+                    bias=bias,
+                    dilation=dilation,
+                )
+
+            def forward(self, x):
+                return self.conv_transpose(x)
+
+        inputs = [torch.randn(1, 3, 32, 32, 32)]
+        self.run_test(TestModule(), inputs, expected_ops={acc_ops.conv_transpose3d})
+
+    def test_conv_transpose3d_with_dynamic_shape(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv_transpose = torch.nn.ConvTranspose3d(3, 6, 1)
+
+            def forward(self, x):
+                return self.conv_transpose(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, 3, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 3, 1, 1, 1), (1, 3, 4, 4, 4), (8, 3, 32, 32, 32))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={acc_ops.conv_transpose3d}
+        )
+
 if __name__ == "__main__":
     run_tests()

test/tracer/test_acc_tracer.py

@@ -449,6 +449,50 @@ def forward(self, a: torch.Tensor) -> torch.Tensor:
 
         self.assertTrue(torch.equal(m(input), traced(input)))
 
+    def test_conv_transpose3d(self):
+        """
+        Test that a conv_transpose3d is traced as expected.
+        """
+
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = nn.ConvTranspose3d(8, 7, 3, stride=2)
+
+            def forward(self, a: torch.Tensor) -> torch.Tensor:
+                return self.conv(a)
+
+        m = TestModule()
+        input = torch.randn(3, 8, 8, 10, 10)
+        traced = acc_tracer.trace(m, [input])
+
+        ph = weight_attr = bias_attr = conv = None
+        for node in traced.graph.nodes:
+            if node.op == "placeholder":
+                self.assertEqual(str(node.target), "a")
+                ph = node
+            elif node.op == "get_attr" and node.target == "conv.weight":
+                weight_attr = node
+            elif node.op == "get_attr" and node.target == "conv.bias":
+                bias_attr = node
+            elif node.op == "call_function":
+                self.assertEqual(node.target, acc_ops.conv_transpose3d)
+                self.assertEqual(node.kwargs["input"], ph)
+                self.assertEqual(node.kwargs["weight"], weight_attr)
+                self.assertEqual(node.kwargs["bias"], bias_attr)
+                self.assertEqual(node.kwargs["stride"], (2, 2, 2))
+                self.assertEqual(node.kwargs["padding"], (0, 0, 0))
+                self.assertEqual(node.kwargs["output_padding"], (0, 0, 0))
+                self.assertEqual(node.kwargs["dilation"], (1, 1, 1))
+                self.assertEqual(node.kwargs["groups"], 1)
+                conv = node
+            elif node.op == "output":
+                self.assertEqual(conv, node.args[0])
+            else:
+                self.fail(f"Unexpected node: {node.format_node()}")
+
+        self.assertTrue(torch.equal(m(input), traced(input)))
+
     def test_embedding_bag(self):
         """
         Test that an embedding_bag is traced as expected.
@@ -2283,6 +2327,7 @@ def test_all_acc_ops_registered(self):
                 acc_ops.conv2d,
                 acc_ops.conv3d,
                 acc_ops.conv_transpose2d,
+                acc_ops.conv_transpose3d,
                 acc_ops.batch_norm,
                 acc_ops.embedding_bag,
                 acc_ops.embedding_bag_byte_rowwise_offsets,

tracer/acc_tracer/acc_ops.py

@@ -1446,6 +1446,21 @@ def conv_transpose2d(*, input, weight, bias, stride, padding, output_padding, gr
     )
 
 
+@register_acc_op_mapping(op_and_target=("call_function", torch.nn.functional.conv_transpose3d))
+@register_acc_op
+def conv_transpose3d(*, input, weight, bias, stride, padding, output_padding, groups, dilation):
+    return nn.functional.conv_transpose3d(
+        input=input,
+        weight=weight,
+        bias=bias,
+        stride=stride,
+        padding=padding,
+        output_padding=output_padding,
+        groups=groups,
+        dilation=dilation,
+    )
+
+
 @register_acc_op_mapping(op_and_target=("call_function", nn.functional.batch_norm))
 @register_acc_op
 def batch_norm(