Optim-wip: Add JIT support to all transforms & some image parameterizations (#821)

* Add JIT support to most transforms

* Additional improvements

* JIT support for `center_crop`.
* Improve some transform tests.
* Fix `RandomCrop` transform bug.

* Fix Mypy bug

* Interpolation based RandomScale & Other Improvements

* Replace Affine `RandomScale` with Interpolation based variant. Renamed old variant to `RandomScaleAffine`.
* `CenterCrop` & `center_crop` now use padding if the crop size is larger than the input dimensions.
* Add distributions support to both versions of `RandomScale`.
* Improve transform tests.

* NumSeqOrTensorType -> NumSeqOrTensorOrProbDistType

* Add `torch.distributions.distribution.Distribution` to `NumSeqOrTensorType` type hint.

* Add TransformationRobustness transform& fix bug

* Added `TransformationRobustness()` transform.
* Fixed bug with `center_crop` padding code, and added related tests to `center_crop` & `CenterCrop`.

* Fix center crop JIT tests

* Add asserts & more tests for RandomScale transforms

* Add JIT support for ToRGB, NaturalImage, & FFTImage

* Add JIT support `NaturalImage`, `FFTImage`, & `PixelImage`.
* Added proper JIT support for `ToRGB`.
* Improved `NaturalImage` & `FFTImage` tests, and test coverage.

* Add ImageParameterization Instance support for NaturalImage

* Added `ImageParameterization` instance support for `NaturalImage`. This improvement should make it easier to use parameterization enhancements like SharedImage, and will be helpful for custom parameterizations that don't use the standard input variable set (size, channels, batch, & init).
* Added asserts to verify `NaturalImage` parameterization inputs are instances or types of `ImageParameterization`.

* Support ToRGB with no named dimensions

This should make it easier to work with the ToRGB module as many PyTorch functions still don't work with named dimensions yet.

* Allow more than 4 channels in ToRGB

* The maximum of 4 channels isn't required as we ignore all channels after 3.

* Add assert check to `RandomScale`'s mode variable

The `linear` mode only supports 3D inputs, and `trilinear` only supports 5D inputs. RandomScale only uses 4D inputs, so only `nearest`, `bilinear`, `bicubic`, & `area` are supported.

* Change assert to check for unsupported RandomScale mode options

* Change `RandomRotation` type hint & add `RandomRotation` to `TransformationRobustness`

* Change `RandomRotation` type hint from `NumSeqOrTensorType` to `NumSeqOrTensorOrProbDistType`.
* Uncomment `RandomRotation` from `TransformationRobustness` & tests.

Author: ProGamerGov

captum/optim/_param/image/images.py

@@ -138,6 +138,8 @@ class FFTImage(ImageParameterization):
     Parameterize an image using inverse real 2D FFT
     """
 
+    __constants__ = ["size", "_supports_is_scripting"]
+
     def __init__(
         self,
         size: Tuple[int, int] = None,
@@ -197,6 +199,9 @@ def __init__(
         self.register_buffer("spectrum_scale", spectrum_scale)
         self.fourier_coeffs = nn.Parameter(fourier_coeffs)
 
+        # Check & store whether or not we can use torch.jit.is_scripting()
+        self._supports_is_scripting = torch.__version__ >= "1.6.0"
+
     def rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
         """
         Computes 2D spectrum frequencies.
@@ -214,6 +219,12 @@ def rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
         fx = self.torch_fftfreq(width)[: width // 2 + 1]
         return torch.sqrt((fx * fx) + (fy * fy))
 
+    @torch.jit.export
+    def torch_irfftn(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dtype != torch.complex64:
+            x = torch.view_as_complex(x)
+        return torch.fft.irfftn(x, s=self.size)  # type: ignore
+
     def get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         """
         Support older versions of PyTorch. This function ensures that the same FFT
@@ -226,26 +237,24 @@ def get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         """
 
         if TORCH_VERSION >= "1.7.0":
-            import torch.fft
+            if TORCH_VERSION < "1.8.0":
+                global torch
+                import torch.fft
 
             def torch_rfft(x: torch.Tensor) -> torch.Tensor:
                 return torch.view_as_real(torch.fft.rfftn(x, s=self.size))
 
-            def torch_irfft(x: torch.Tensor) -> torch.Tensor:
-                if type(x) is not torch.complex64:
-                    x = torch.view_as_complex(x)
-                return torch.fft.irfftn(x, s=self.size)  # type: ignore
+            torch_irfftn = self.torch_irfftn
 
             def torch_fftfreq(v: int, d: float = 1.0) -> torch.Tensor:
                 return torch.fft.fftfreq(v, d)
 
         else:
-            import torch
 
             def torch_rfft(x: torch.Tensor) -> torch.Tensor:
                 return torch.rfft(x, signal_ndim=2)
 
-            def torch_irfft(x: torch.Tensor) -> torch.Tensor:
+            def torch_irfftn(x: torch.Tensor) -> torch.Tensor:
                 return torch.irfft(x, signal_ndim=2)[
                     :, :, : self.size[0], : self.size[1]
                 ]
@@ -258,7 +267,7 @@ def torch_fftfreq(v: int, d: float = 1.0) -> torch.Tensor:
                 results[s:] = torch.arange(-(v // 2), 0)
                 return results * (1.0 / (v * d))
 
-        return torch_rfft, torch_irfft, torch_fftfreq
+        return torch_rfft, torch_irfftn, torch_fftfreq
 
     def forward(self) -> torch.Tensor:
         """
@@ -268,6 +277,9 @@ def forward(self) -> torch.Tensor:
 
         scaled_spectrum = self.fourier_coeffs * self.spectrum_scale
         output = self.torch_irfft(scaled_spectrum)
+        if self._supports_is_scripting:
+            if torch.jit.is_scripting():
+                return output
         return output.refine_names("B", "C", "H", "W")
 
 
@@ -276,6 +288,8 @@ class PixelImage(ImageParameterization):
     Parameterize a simple pixel image tensor that requires no additional transforms.
     """
 
+    __constants__ = ["_supports_is_scripting"]
+
     def __init__(
         self,
         size: Tuple[int, int] = None,
@@ -309,7 +323,13 @@ def __init__(
             f"input has {init.shape[1]} channels."
         self.image = nn.Parameter(init)
 
+        # Check & store whether or not we can use torch.jit.is_scripting()
+        self._supports_is_scripting = torch.__version__ >= "1.6.0"
+
     def forward(self) -> torch.Tensor:
+        if self._supports_is_scripting:
+            if torch.jit.is_scripting():
+                return self.image
         return self.image.refine_names("B", "C", "H", "W")
 
 
@@ -600,7 +620,7 @@ def __init__(
                 nn.Parameter tensor, or stacking init images.
                 Default: 1
             parameterization (ImageParameterization, optional): An image
-                parameterization class.
+                parameterization class, or instance of an image parameterization class.
                 Default: FFTImage
             squash_func (Callable[[torch.Tensor], torch.Tensor]], optional): The squash
                 function to use after color recorrelation. A funtion or lambda function.
@@ -612,8 +632,14 @@ def __init__(
                 Default: True
         """
         super().__init__()
+        if not isinstance(parameterization, ImageParameterization):
+            # Verify uninitialized class is correct type
+            assert issubclass(parameterization, ImageParameterization)
+        else:
+            assert isinstance(parameterization, ImageParameterization)
+
         self.decorrelate = decorrelation_module
-        if init is not None:
+        if init is not None and not isinstance(parameterization, ImageParameterization):
             assert init.dim() == 3 or init.dim() == 4
             if decorrelate_init and self.decorrelate is not None:
                 init = (
@@ -622,27 +648,42 @@ def __init__(
                     else init.refine_names("C", "H", "W")
                 )
                 init = self.decorrelate(init, inverse=True).rename(None)
+
             if squash_func is None:
+                squash_func = self._clamp_image
 
-                def squash_func(x: torch.Tensor) -> torch.Tensor:
-                    return x.clamp(0, 1)
+        self.squash_func = torch.sigmoid if squash_func is None else squash_func
+        if not isinstance(parameterization, ImageParameterization):
+            parameterization = parameterization(
+                size=size, channels=channels, batch=batch, init=init
+            )
+        self.parameterization = parameterization
 
-        else:
-            if squash_func is None:
+    @torch.jit.export
+    def _clamp_image(self, x: torch.Tensor) -> torch.Tensor:
+        """JIT supported squash function."""
+        return x.clamp(0, 1)
 
-                squash_func = torch.sigmoid
+    @torch.jit.ignore
+    def _to_image_tensor(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Wrap ImageTensor in torch.jit.ignore for JIT support.
 
-        self.squash_func = squash_func
-        self.parameterization = parameterization(
-            size=size, channels=channels, batch=batch, init=init
-        )
+        Args:
+
+            x (torch.tensor): An input tensor.
+
+        Returns:
+            x (ImageTensor): An instance of ImageTensor with the input tensor.
+        """
+        return ImageTensor(x)
 
     def forward(self) -> torch.Tensor:
         image = self.parameterization()
         if self.decorrelate is not None:
             image = self.decorrelate(image)
         image = image.rename(None)  # TODO: the world is not yet ready
-        return ImageTensor(self.squash_func(image))
+        return self._to_image_tensor(self.squash_func(image))
 
 
 __all__ = [