feat: support native_dropout dynamo converter (#2931)

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -3243,3 +3243,36 @@ def aten_ops_index_select(
         args[1],
         args[2],
     )
+
+
+def dropout_inference_validator(node: Node) -> bool:
+    train_mode = args_bounds_check(node.args, 2, None)
+    if train_mode is False:
+        return True
+    else:  # train_mode is True or None
+        _LOGGER.debug(
+            "Currently only inference mode is supported for dropout operation."
+        )
+        return False
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.native_dropout.default,
+    capability_validator=dropout_inference_validator,
+)
+def aten_ops_native_dropout(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unary.native_dropout(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args_bounds_check(args, 2, None),
+    )

py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py

@@ -571,3 +571,20 @@ def isnan(
     )
 
     return nan_values_mask
+
+
+def native_dropout(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_val: Union[TRTTensor, torch.Tensor, np.ndarray],
+    p: float,
+    train: Optional[bool] = False,
+) -> TRTTensor:
+    if train is False:
+        identity_layer = ctx.net.add_identity(input_val)
+        set_layer_name(identity_layer, target, f"{name}_input", source_ir)
+        mask = np.ones(input_val.shape, dtype=bool)
+        mask = get_trt_tensor(ctx, mask, f"{name}_mask")
+        return identity_layer.get_output(0), mask

tests/py/dynamo/conversion/test_dropout_aten.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestDropOutConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((10,), 0, False),
+            ((1, 3), 0.3, False),
+            ((2, 2, 2), 0.5),
+            ((2, 2, 2, 2), 1),
+        ]
+    )
+    def test_native_dropout(self, input_shape, p, train=False):
+        class NativeDropout(nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.native_dropout.default(input, p, train)
+
+        inputs = [torch.randn(input_shape)]
+        self.run_test(
+            NativeDropout(),
+            inputs,
+        )
+
+    @parameterized.expand(
+        [
+            (
+                torch.randn(
+                    10,
+                ),
+                0,
+                False,
+            ),
+            (torch.randn(1, 3), 0.3, False),
+            (torch.randn(2, 2, 2), 0.5),
+            (torch.randn(2, 2, 2, 2), 1),
+        ]
+    )
+    def test_native_dropout_pytorch(self, input, p, train=False):
+        class NativeDropout(nn.Module):
+            def forward(self):
+                return torch.ops.aten.native_dropout.default(input, p, train)
+
+        inputs = []
+        self.run_test(
+            NativeDropout(),
+            inputs,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()