refactor: Centralizing sigmoid implementation

Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -3156,17 +3156,8 @@ def acc_ops_sigmoid(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"Sigmoid received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
 
-    return add_activation_layer(
-        network, input_val, trt.ActivationType.SIGMOID, target, name
-    )
+    return activation.convert_sigmoid(network, target, kwargs, name, SourceIR.ACC)
 
 
 @tensorrt_converter(acc_ops.permute)

py/torch_tensorrt/fx/converters/activation.py

@@ -482,3 +482,19 @@ def aten_ops_sym_size(
     )
     set_layer_name(slice_layer, target, "_slice_layer")
     return slice_layer.get_output(0)
+
+
+@tensorrt_converter(torch.ops.aten.sigmoid.default)
+def aten_ops_sigmoid(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+
+    kwargs_new = {
+        "input": args[0],
+    }
+
+    return activation.convert_sigmoid(network, target, kwargs_new, name, SourceIR.ATEN)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -85,3 +85,31 @@ def relu_dyn_range_fn(dyn_range):
     return convert_activation(
         network, input_val, operation_type, target, name, relu_dyn_range_fn, source_ir
     )
+
+
+def convert_sigmoid(
+    network: TRTNetwork,
+    target: Target,
+    kwargs: Dict[str, Any],
+    name: str,
+    source_ir: SourceIR = SourceIR.UNKNOWN,
+):
+    input_val = kwargs["input"]
+    operation_type = trt.ActivationType.SIGMOID
+
+    def sigmoid_dyn_range_fn(dyn_range):
+        def sigmoid_fn(x):
+            # TODO: Can this just call torch.nn.functional.sigmoid?
+            return 1 / (1 + np.exp(-x))
+
+        return sigmoid_fn(dyn_range[0]), sigmoid_fn(dyn_range[1])
+
+    return convert_activation(
+        network,
+        input_val,
+        operation_type,
+        target,
+        name,
+        sigmoid_dyn_range_fn,
+        source_ir,
+    )

py/torch_tensorrt/fx/converters/impl/activation.py

@@ -22,3 +22,17 @@ def relu(network, submod, args, kwargs, layer_name):
         name=layer_name,
         source_ir=SourceIR.NN,
     )
+
+
+@tensorrt_converter(torch.nn.modules.activation.Sigmoid)
+def sigmoid(network, submod, args, kwargs, layer_name):
+    # args/kwargs should have already been normalized to kwargs
+    assert len(args) == 0
+
+    activation.convert_sigmoid(
+        network=network,
+        target="torch.nn.modules.activation.Sigmoid",
+        kwargs=kwargs,
+        name=layer_name,
+        source_ir=SourceIR.NN,
+    )

py/torch_tensorrt/fx/converters/nn_ops_converters.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()