[fbsync] allow nn.ModuleList in RandomApply (#7197)

Reviewed By: vmoens

Differential Revision: D44416259

fbshipit-source-id: cbd7c731f37f13b3bdb4d117225910fa019eac60

test/test_prototype_transforms_consistency.py

@@ -23,6 +23,7 @@
     make_label,
     make_segmentation_mask,
 )
+from torch import nn
 from torchvision import transforms as legacy_transforms
 from torchvision._utils import sequence_to_str
 from torchvision.prototype import datapoints, transforms as prototype_transforms
@@ -761,19 +762,24 @@ def test_compose(self):
         check_call_consistency(prototype_transform, legacy_transform)
 
     @pytest.mark.parametrize("p", [0, 0.1, 0.5, 0.9, 1])
-    def test_random_apply(self, p):
+    @pytest.mark.parametrize("sequence_type", [list, nn.ModuleList])
+    def test_random_apply(self, p, sequence_type):
         prototype_transform = prototype_transforms.RandomApply(
-            [
-                prototype_transforms.Resize(256),
-                prototype_transforms.CenterCrop(224),
-            ],
+            sequence_type(
+                [
+                    prototype_transforms.Resize(256),
+                    prototype_transforms.CenterCrop(224),
+                ]
+            ),
             p=p,
         )
         legacy_transform = legacy_transforms.RandomApply(
-            [
-                legacy_transforms.Resize(256),
-                legacy_transforms.CenterCrop(224),
-            ],
+            sequence_type(
+                [
+                    legacy_transforms.Resize(256),
+                    legacy_transforms.CenterCrop(224),
+                ]
+            ),
             p=p,
         )
 

torchvision/prototype/transforms/_container.py

@@ -1,7 +1,9 @@
 import warnings
-from typing import Any, Callable, List, Optional, Sequence
+from typing import Any, Callable, List, Optional, Sequence, Union
 
 import torch
+
+from torch import nn
 from torchvision.prototype.transforms import Transform
 
 
@@ -25,9 +27,13 @@ def extra_repr(self) -> str:
         return "\n".join(format_string)
 
 
-class RandomApply(Compose):
-    def __init__(self, transforms: Sequence[Callable], p: float = 0.5) -> None:
-        super().__init__(transforms)
+class RandomApply(Transform):
+    def __init__(self, transforms: Union[Sequence[Callable], nn.ModuleList], p: float = 0.5) -> None:
+        super().__init__()
+
+        if not isinstance(transforms, (Sequence, nn.ModuleList)):
+            raise TypeError("Argument transforms should be a sequence of callables or a `nn.ModuleList`")
+        self.transforms = transforms
 
         if not (0.0 <= p <= 1.0):
             raise ValueError("`p` should be a floating point value in the interval [0.0, 1.0].")
@@ -39,7 +45,15 @@ def forward(self, *inputs: Any) -> Any:
         if torch.rand(1) >= self.p:
             return sample
 
-        return super().forward(sample)
+        for transform in self.transforms:
+            sample = transform(sample)
+        return sample
+
+    def extra_repr(self) -> str:
+        format_string = []
+        for t in self.transforms:
+            format_string.append(f"    {t}")
+        return "\n".join(format_string)
 
 
 class RandomChoice(Transform):