Specialize to TA to TAwide

Author: datumbox

references/classification/presets.py

@@ -9,8 +9,8 @@ def __init__(self, crop_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.2
         if hflip_prob > 0:
             trans.append(transforms.RandomHorizontalFlip(hflip_prob))
         if auto_augment_policy is not None:
-            if auto_augment_policy == autoaugment.AugmentationSpace.TA_WIDE.value:
-                trans.append(autoaugment.TrivialAugment())
+            if auto_augment_policy == "ta_wide":
+                trans.append(autoaugment.TrivialAugmentWide())
             else:
                 aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
                 trans.append(autoaugment.AutoAugment(policy=aa_policy))

test/test_transforms.py

@@ -1490,14 +1490,12 @@ def test_autoaugment(policy, fill):
     transform.__repr__()
 
 
-@pytest.mark.parametrize('augmentation_space', [space for space in transforms.AugmentationSpace])
 @pytest.mark.parametrize('fill', [None, 85, (128, 128, 128)])
 @pytest.mark.parametrize('num_magnitude_bins', [10, 13, 30])
-def test_trivialaugment(augmentation_space, fill, num_magnitude_bins):
+def test_trivialaugmentwide(fill, num_magnitude_bins):
     random.seed(42)
     img = Image.open(GRACE_HOPPER)
-    transform = transforms.TrivialAugment(augmentation_space=augmentation_space,
-                                          fill=fill, num_magnitude_bins=num_magnitude_bins)
+    transform = transforms.TrivialAugmentWide(fill=fill, num_magnitude_bins=num_magnitude_bins)
     for _ in range(100):
         img = transform(img)
     transform.__repr__()

test/test_transforms_tensor.py

@@ -540,22 +540,20 @@ def test_autoaugment_save(tmpdir):
 
 
 @pytest.mark.parametrize('device', cpu_and_gpu())
-@pytest.mark.parametrize('augmentation_space', [space for space in T.AugmentationSpace])
 @pytest.mark.parametrize('fill', [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1])
-def test_trivialaugment(device, augmentation_space, fill):
+def test_trivialaugmentwide(device, fill):
     tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
     batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
 
-    s_transform = None
-    transform = T.TrivialAugment(augmentation_space=augmentation_space, fill=fill)
+    transform = T.TrivialAugmentWide(fill=fill)
     s_transform = torch.jit.script(transform)
     for _ in range(25):
         _test_transform_vs_scripted(transform, s_transform, tensor)
         _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
 
 
-def test_trivialaugment_save(tmpdir):
-    transform = T.TrivialAugment()
+def test_trivialaugmentwide_save(tmpdir):
+    transform = T.TrivialAugmentWide()
     s_transform = torch.jit.script(transform)
     s_transform.save(os.path.join(tmpdir, "t_autoaugment.pt"))
 

torchvision/transforms/autoaugment.py

@@ -7,7 +7,7 @@
 
 from . import functional as F, InterpolationMode
 
-__all__ = ["AutoAugmentPolicy", "AutoAugment", "AugmentationSpace", "TrivialAugment"]
+__all__ = ["AutoAugmentPolicy", "AutoAugment", "TrivialAugmentWide"]
 
 
 def _apply_op(img: Tensor, op_name: str, magnitude: float,
@@ -178,8 +178,7 @@ def _get_transforms(
         else:
             raise ValueError("The provided policy {} is not recognized.".format(policy))
 
-    @staticmethod
-    def _get_magnitudes(num_bins: int, image_size: List[int]) -> Dict[str, Tuple[Tensor, bool]]:
+    def _get_magnitudes(self, num_bins: int, image_size: List[int]) -> Dict[str, Tuple[Tensor, bool]]:
         return {
             # name: (magnitudes, signed)
             "ShearX": (torch.linspace(0.0, 0.3, num_bins), True),
@@ -243,24 +242,14 @@ def __repr__(self) -> str:
         return self.__class__.__name__ + '(policy={}, fill={})'.format(self.policy, self.fill)
 
 
-class AugmentationSpace(Enum):
-    """The augmentation space to use.
-    Available spaces are `AA` for AutoAugment and `TA_WIDE` for the TrivialAugment.
-    """
-    AA = "aa"
-    TA_WIDE = "ta_wide"
-
-
-class TrivialAugment(torch.nn.Module):
-    r"""Dataset-independent data-augmentation with TrivialAugment, as described in
+class TrivialAugmentWide(torch.nn.Module):
+    r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
     `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`.
         If the image is torch Tensor, it should be of type torch.uint8, and it is expected
         to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
         If img is PIL Image, it is expected to be in mode "L" or "RGB".
 
         Args:
-            augmentation_space (AugmentationSpace): Desired augmentation space enum defined by
-                :class:`torchvision.transforms.autoaugment.AugmentationSpace`. Default is ``AugmentationSpace.TA_WIDE``.
             num_magnitude_bins (int): The number of different magnitude values.
             interpolation (InterpolationMode): Desired interpolation enum defined by
                 :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
@@ -269,17 +258,14 @@ class TrivialAugment(torch.nn.Module):
                 image. If given a number, the value is used for all bands respectively.
         """
 
-    def __init__(self, augmentation_space: AugmentationSpace = AugmentationSpace.TA_WIDE, num_magnitude_bins: int = 30,
-                 interpolation: InterpolationMode = InterpolationMode.NEAREST,
+    def __init__(self, num_magnitude_bins: int = 30, interpolation: InterpolationMode = InterpolationMode.NEAREST,
                  fill: Optional[List[float]] = None) -> None:
         super().__init__()
-        self.augmentation_space = augmentation_space
         self.num_magnitude_bins = num_magnitude_bins
         self.interpolation = interpolation
         self.fill = fill
 
-    @staticmethod
-    def _get_magnitudes(num_bins: int) -> Dict[str, Tuple[Tensor, bool]]:
+    def _get_magnitudes(self, num_bins: int) -> Dict[str, Tuple[Tensor, bool]]:
         return {
             # name: (magnitudes, signed)
             "ShearX": (torch.linspace(0.0, 0.99, num_bins), True),
@@ -303,7 +289,7 @@ def forward(self, img: Tensor):
             img (PIL Image or Tensor): Image to be transformed.
 
         Returns:
-            PIL Image or Tensor: TrivialAugmented image.
+            PIL Image or Tensor: Transformed image.
         """
         fill = self.fill
         if isinstance(img, Tensor):
@@ -312,12 +298,7 @@ def forward(self, img: Tensor):
             elif fill is not None:
                 fill = [float(f) for f in fill]
 
-        if self.augmentation_space == AugmentationSpace.AA:
-            op_meta = AutoAugment._get_magnitudes(self.num_magnitude_bins, F.get_image_size(img))
-        elif self.augmentation_space == AugmentationSpace.TA_WIDE:
-            op_meta = self._get_magnitudes(self.num_magnitude_bins)
-        else:
-            raise ValueError(f"Provided augmentation_space arguments {self.augmentation_space} not available.")
+        op_meta = self._get_magnitudes(self.num_magnitude_bins)
         op_index = int(torch.randint(len(op_meta), (1,)).item())
         op_name = list(op_meta.keys())[op_index]
         magnitudes, signed = op_meta[op_name]
@@ -330,8 +311,7 @@ def forward(self, img: Tensor):
 
     def __repr__(self) -> str:
         s = self.__class__.__name__ + '('
-        s += 'augmentation_space={augmentation_space}'
-        s += ', num_magnitude_bins={num_magnitude_bins}'
+        s += 'num_magnitude_bins={num_magnitude_bins}'
         s += ', interpolation={interpolation}'
         s += ', fill={fill}'
         s += ')'