correcting selu, hard_tanh ops, adding tests for sigmoid, selu, elu and hard_tanh

Author: apbose

py/torch_tensorrt/fx/converters/activation.py

@@ -184,25 +184,14 @@ def add_sigmoid(network, target, kwargs, name):
 
 def add_hard_tanh(network, target, kwargs, name):
     input_val = kwargs["input"]
-    operation_type = trt.ActivationType.TANH
-    return add_activation_layer(network, input_val, operation_type, target, name)
-
-
-def add_sigmoid(network, target, kwargs, name):
-    input_val = kwargs["input"]
-
+    alpha = kwargs["min_val"]
+    beta = kwargs["max_val"]
     if not isinstance(input_val, TRTTensor):
         raise RuntimeError(
-            f"Hard sigmoid received input {input_val} that is not part "
+            f"hardtanh received input {input_val} that is not part "
             "of the TensorRT region!"
         )
-
+    operation_type = trt.ActivationType.CLIP
     return add_activation_layer(
-        network,
-        input_val,
-        trt.ActivationType.HARD_SIGMOID,
-        target,
-        name,
-        alpha=1 / 6,
-        beta=0.5,
+        network, input_val, operation_type, target, name, alpha, beta
     )

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -503,26 +503,19 @@ def aten_ops_elu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-    }
+    if len(args) > 2:
+        kwargs_new = {
+            "input": args[0],
+        }
+        return add_selu(network, target, kwargs_new, name)
+    else:
+        kwargs_new = {
+            "input": args[0],
+            "alpha": args[1],
+        }
     return add_elu(network, target, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.selu.default)
-def aten_ops_selu(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-    }
-    return add_selu(network, target, kwargs_new, name)
-
-
 @tensorrt_converter(torch.ops.aten.gelu.default)
 def aten_ops_gelu(
     network: TRTNetwork,
@@ -551,7 +544,7 @@ def aten_ops_tanh(
     return add_tanh(network, target, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.sigmoid.default)
+@tensorrt_converter(torch.ops.aten.hardtanh.default)
 def aten_ops_hard_tanh(
     network: TRTNetwork,
     target: Target,
@@ -561,6 +554,8 @@ def aten_ops_hard_tanh(
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     kwargs_new = {
         "input": args[0],
+        "min_val": args[1],
+        "max_val": args[2],
     }
     return add_hard_tanh(network, target, kwargs_new, name)
 

py/torch_tensorrt/fx/test/converters/aten_op/test_elu_aten.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestELUConverter(DispatchTestCase):
+    def test_elu(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.elu(x)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.elu.default})
+
+    def test_elu_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.elu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.elu.default}
+        )
+
+    def test_elu_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.elu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.elu.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

py/torch_tensorrt/fx/test/converters/aten_op/test_hard_tanh_aten.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestHardTanHConverter(DispatchTestCase):
+    def test_hardtanh(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.hardtanh(x)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(
+            TestModule(), inputs, expected_ops={torch.ops.aten.hardtanh.default}
+        )
+
+    def test_hardtanh_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.hardtanh(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.hardtanh.default}
+        )
+
+    def test_hardtanh_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.hardtanh(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.hardtanh.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

py/torch_tensorrt/fx/test/converters/aten_op/test_leaky_relu_aten.py

@@ -4,7 +4,7 @@
 from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
 
 
-class TestReLUConverter(DispatchTestCase):
+class TestLeakyReLUConverter(DispatchTestCase):
     def test_leaky_relu(self):
         class TestModule(nn.Module):
             def forward(self, x):

py/torch_tensorrt/fx/test/converters/aten_op/test_selu_aten.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestSeLUConverter(DispatchTestCase):
+    def test_selu(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.selu(x)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.elu.default})
+
+    def test_selu_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.selu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.elu.default}
+        )
+
+    def test_selu_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.selu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.elu.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

py/torch_tensorrt/fx/test/converters/aten_op/test_sigmoid_aten.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestSigmoidConverter(DispatchTestCase):
+    def test_sigmoid(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.sigmoid(x)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(
+            TestModule(), inputs, expected_ops={torch.ops.aten.sigmoid.default}
+        )
+
+    def test_sigmoid_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.sigmoid(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.sigmoid.default}
+        )
+
+    def test_sigmoid_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.sigmoid(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.sigmoid.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()