Add asserts & more tests for RandomScale transforms

Author: ProGamerGov

captum/optim/_param/image/transforms.py

@@ -358,10 +358,10 @@ class RandomScale(nn.Module):
         "scale",
         "mode",
         "align_corners",
-        "has_align_corners",
+        "_has_align_corners",
         "recompute_scale_factor",
-        "has_recompute_scale_factor",
-        "is_distribution",
+        "_has_recompute_scale_factor",
+        "_is_distribution",
     ]
 
     def __init__(
@@ -390,8 +390,9 @@ def __init__(
         super().__init__()
         if isinstance(scale, torch.distributions.distribution.Distribution):
             # Distributions are not supported by TorchScript / JIT yet
+            assert scale.batch_shape == torch.Size([])
             self.scale_distribution = scale
-            self.is_distribution = True
+            self._is_distribution = True
             self.scale = []
         else:
             assert hasattr(scale, "__iter__")
@@ -400,33 +401,12 @@ def __init__(
                 scale = scale.tolist()
             assert len(scale) > 0
             self.scale = [float(s) for s in scale]
-            self.is_distribution = False
+            self._is_distribution = False
         self.mode = mode
         self.align_corners = align_corners if mode not in ["nearest", "area"] else None
         self.recompute_scale_factor = recompute_scale_factor
-        self.has_align_corners = torch.__version__ >= "1.3.0"
-        self.has_recompute_scale_factor = torch.__version__ >= "1.6.0"
-
-    def _get_scale_mat(
-        self,
-        m: float,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> torch.Tensor:
-        """
-        Create a scale matrix tensor.
-
-        Args:
-
-            m (float): The scale value to use.
-
-        Returns:
-            **scale_mat** (torch.Tensor): A scale matrix.
-        """
-        scale_mat = torch.tensor(
-            [[m, 0.0, 0.0], [0.0, m, 0.0]], device=device, dtype=dtype
-        )
-        return scale_mat
+        self._has_align_corners = torch.__version__ >= "1.3.0"
+        self._has_recompute_scale_factor = torch.__version__ >= "1.6.0"
 
     def _scale_tensor(self, x: torch.Tensor, scale: float) -> torch.Tensor:
         """
@@ -440,8 +420,8 @@ def _scale_tensor(self, x: torch.Tensor, scale: float) -> torch.Tensor:
         Returns:
             **x** (torch.Tensor): A scaled NCHW image tensor.
         """
-        if self.has_align_corners:
-            if self.has_recompute_scale_factor:
+        if self._has_align_corners:
+            if self._has_recompute_scale_factor:
                 x = F.interpolate(
                     x,
                     scale_factor=scale,
@@ -472,8 +452,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             **x** (torch.Tensor): A randomly scaled NCHW image *tensor*.
         """
         assert x.dim() == 4
-        if self.is_distribution:
-            scale = self.scale_distribution.sample().item()
+        if self._is_distribution:
+            scale = float(self.scale_distribution.sample().item())
         else:
             n = int(
                 torch.randint(
@@ -508,8 +488,8 @@ class RandomScaleAffine(nn.Module):
         "mode",
         "padding_mode",
         "align_corners",
-        "has_align_corners",
-        "is_distribution",
+        "_has_align_corners",
+        "_is_distribution",
     ]
 
     def __init__(
@@ -539,8 +519,9 @@ def __init__(
         super().__init__()
         if isinstance(scale, torch.distributions.distribution.Distribution):
             # Distributions are not supported by TorchScript / JIT yet
+            assert scale.batch_shape == torch.Size([])
             self.scale_distribution = scale
-            self.is_distribution = True
+            self._is_distribution = True
             self.scale = []
         else:
             assert hasattr(scale, "__iter__")
@@ -549,11 +530,11 @@ def __init__(
                 scale = scale.tolist()
             assert len(scale) > 0
             self.scale = [float(s) for s in scale]
-            self.is_distribution = False
+            self._is_distribution = False
         self.mode = mode
         self.padding_mode = padding_mode
         self.align_corners = align_corners
-        self.has_align_corners = torch.__version__ >= "1.3.0"
+        self._has_align_corners = torch.__version__ >= "1.3.0"
 
     def _get_scale_mat(
         self,
@@ -591,7 +572,7 @@ def _scale_tensor(self, x: torch.Tensor, scale: float) -> torch.Tensor:
         scale_matrix = self._get_scale_mat(scale, x.device, x.dtype)[None, ...].repeat(
             x.shape[0], 1, 1
         )
-        if self.has_align_corners:
+        if self._has_align_corners:
             # Pass align_corners explicitly for torch >= 1.3.0
             grid = F.affine_grid(
                 scale_matrix, x.size(), align_corners=self.align_corners
@@ -620,8 +601,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             **x** (torch.Tensor): A randomly scaled NCHW image *tensor*.
         """
         assert x.dim() == 4
-        if self.is_distribution:
-            scale = self.scale_distribution.sample().item()
+        if self._is_distribution:
+            scale = float(self.scale_distribution.sample().item())
         else:
             n = int(
                 torch.randint(
@@ -1021,7 +1002,7 @@ class TransformationRobustness(nn.Module):
 
     def __init__(
         self,
-        padding_transform: Optional[nn.Module] = None,
+        padding_transform: Optional[nn.Module] = nn.ConstantPad2d(2, value=0.5),
         translate: Optional[Union[int, List[int]]] = [4] * 10,
         scale: Optional[NumSeqOrTensorOrProbDistType] = [
             0.995 ** n for n in range(-5, 80)
@@ -1039,7 +1020,7 @@ def __init__(
 
             padding_transform (nn.Module, optional): A padding module instance. No
                 padding will be applied before transforms if set to None.
-                Default: None
+                Default: nn.ConstantPad2d(2, value=0.5)
             translate (int or list of int, optional): The max horizontal and vertical
                  translation to use for each jitter transform.
                  Default: [4] * 10