More tests & new AugmentedImageParameterization base class

* Added `AugmentedImageParameterization` class to use a base for `SharedImage` and `StackImage`.
* Removed `PixelImage`'s 3 channel assert, as there was no reason for limitation.
* Added tests for `InputParameterization`, `ImageParameterization`, & `AugmentedImageParameterization`.

Author: ProGamerGov

captum/optim/_param/image/images.py

@@ -133,6 +133,10 @@ class ImageParameterization(InputParameterization):
     pass
 
 
+class AugmentedImageParameterization(ImageParameterization):
+    pass
+
+
 class FFTImage(ImageParameterization):
     """
     Parameterize an image using inverse real 2D FFT
@@ -305,8 +309,6 @@ def __init__(
             assert init.dim() == 3 or init.dim() == 4
             if init.dim() == 3:
                 init = init.unsqueeze(0)
-            assert init.shape[1] == 3, "PixelImage init should have 3 channels, "
-            f"input has {init.shape[1]} channels."
         self.image = nn.Parameter(init)
 
     def forward(self) -> torch.Tensor:
@@ -432,10 +434,14 @@ def __init__(self, tensor: torch.Tensor = None) -> None:
         self.tensor = tensor
 
     def forward(self) -> torch.Tensor:
+        """
+        Returns:
+            tensor (torch.Tensor): The tensor stored during initialization.
+        """
         return self.tensor
 
 
-class SharedImage(ImageParameterization):
+class SharedImage(AugmentedImageParameterization):
     """
     Share some image parameters across the batch to increase spatial alignment,
     by using interpolated lower resolution tensors.
@@ -585,6 +591,10 @@ def _interpolate_tensor(
         return x
 
     def forward(self) -> torch.Tensor:
+        """
+        Returns:
+            output (torch.Tensor): An NCHW image parameterization output.
+        """
         image = self.parameterization()
         x = [
             self._interpolate_tensor(
@@ -606,7 +616,7 @@ def forward(self) -> torch.Tensor:
         return output.refine_names("B", "C", "H", "W")
 
 
-class StackImage(ImageParameterization):
+class StackImage(AugmentedImageParameterization):
     """
     Stack multiple NCHW image parameterizations along their batch dimensions.
     """
@@ -691,7 +701,9 @@ def __init__(
         channels: int = 3,
         batch: int = 1,
         init: Optional[torch.Tensor] = None,
-        parameterization: ImageParameterization = FFTImage,
+        parameterization: Union[
+            ImageParameterization, AugmentedImageParameterization
+        ] = FFTImage,
         squash_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
         decorrelation_module: Optional[nn.Module] = ToRGB(transform="klt"),
         decorrelate_init: bool = True,
@@ -758,9 +770,11 @@ def forward(self) -> torch.Tensor:
     "ImageTensor",
     "InputParameterization",
     "ImageParameterization",
+    "AugmentedImageParameterization",
     "FFTImage",
     "PixelImage",
     "LaplacianImage",
     "SharedImage",
+    "StackImage",
     "NaturalImage",
 ]

tests/optim/param/test_images.py

@@ -79,7 +79,7 @@ def test_export_and_open_local_image(self) -> None:
         self.assertTrue(torch.is_tensor(new_tensor))
         assertTensorAlmostEqual(self, image_tensor, new_tensor)
 
-    def test_natural_image_cuda(self) -> None:
+    def test_image_tensor_cuda(self) -> None:
         if not torch.cuda.is_available():
             raise unittest.SkipTest(
                 "Skipping ImageTensor CUDA test due to not supporting CUDA."
@@ -88,7 +88,31 @@ def test_natural_image_cuda(self) -> None:
         self.assertTrue(image_t.is_cuda)
 
 
+class TestInputParameterization(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(issubclass(images.InputParameterization, torch.nn.Module))
+
+
+class TestImageParameterization(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(
+            issubclass(images.ImageParameterization, images.InputParameterization)
+        )
+
+
+class TestAugmentedImageParameterization(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(
+            issubclass(
+                images.AugmentedImageParameterization, images.ImageParameterization
+            )
+        )
+
+
 class TestFFTImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(issubclass(images.FFTImage, images.ImageParameterization))
+
     def test_pytorch_fftfreq(self) -> None:
         image = images.FFTImage((1, 1))
         _, _, fftfreq = image.get_fft_funcs()
@@ -219,6 +243,9 @@ def test_fftimage_forward_init_batch(self) -> None:
 
 
 class TestPixelImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(issubclass(images.PixelImage, images.ImageParameterization))
+
     def test_pixelimage_random(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -251,17 +278,6 @@ def test_pixelimage_init(self) -> None:
         self.assertEqual(image_param.image.size(3), size[1])
         assertTensorAlmostEqual(self, image_param.image, init_tensor, 0)
 
-    def test_pixelimage_init_error(self) -> None:
-        if torch.__version__ <= "1.2.0":
-            raise unittest.SkipTest(
-                "Skipping PixelImage init due to insufficient Torch version."
-            )
-        size = (224, 224)
-        channels = 2
-        init_tensor = torch.randn(channels, *size)
-        with self.assertRaises(AssertionError):
-            images.PixelImage(size=size, channels=channels, init=init_tensor)
-
     def test_pixelimage_random_forward(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -298,6 +314,9 @@ def test_pixelimage_init_forward(self) -> None:
 
 
 class TestLaplacianImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(issubclass(images.LaplacianImage, images.ImageParameterization))
+
     def test_laplacianimage_random_forward(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -326,6 +345,13 @@ def test_laplacianimage_init(self) -> None:
 
 
 class TestSimpleTensorParameterization(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(
+            issubclass(
+                images.SimpleTensorParameterization, images.ImageParameterization
+            )
+        )
+
     def test_simple_tensor_parameterization_no_grad(self) -> None:
         test_input = torch.randn(1, 3, 4, 4)
         image_param = images.SimpleTensorParameterization(test_input)
@@ -337,6 +363,11 @@ def test_simple_tensor_parameterization_no_grad(self) -> None:
         self.assertFalse(image_param.tensor.requires_grad)
 
     def test_simple_tensor_parameterization_jit_module_no_grad(self) -> None:
+        if torch.__version__ <= "1.8.0":
+            raise unittest.SkipTest(
+                "Skipping SimpleTensorParameterization JIT module test due to"
+                + "  insufficient Torch version."
+            )
         test_input = torch.randn(1, 3, 4, 4)
         image_param = images.SimpleTensorParameterization(test_input)
         jit_image_param = torch.jit.script(image_param)
@@ -356,6 +387,11 @@ def test_simple_tensor_parameterization_with_grad(self) -> None:
         self.assertTrue(image_param.tensor.requires_grad)
 
     def test_simple_tensor_parameterization_jit_module_with_grad(self) -> None:
+        if torch.__version__ <= "1.8.0":
+            raise unittest.SkipTest(
+                "Skipping SimpleTensorParameterization JIT module test due to"
+                + "  insufficient Torch version."
+            )
         test_input = torch.nn.Parameter(torch.randn(1, 3, 4, 4))
         image_param = images.SimpleTensorParameterization(test_input)
         jit_image_param = torch.jit.script(image_param)
@@ -383,6 +419,11 @@ def test_simple_tensor_parameterization_cuda(self) -> None:
 
 
 class TestSharedImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(
+            issubclass(images.SharedImage, images.AugmentedImageParameterization)
+        )
+
     def test_sharedimage_get_offset_single_number(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -726,6 +767,11 @@ def test_sharedimage_multiple_shapes_diff_len_forward_jit_module(self) -> None:
 
 
 class TestStackImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(
+            issubclass(images.StackImage, images.AugmentedImageParameterization)
+        )
+
     def test_stackimage_init(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -940,6 +986,9 @@ def test_stackimage_forward_multi_device_cpu_gpu(self) -> None:
 
 
 class TestNaturalImage(BaseTest):
+    def test_subclass(self) -> None:
+        self.assertTrue(issubclass(images.NaturalImage, images.ImageParameterization))
+
     def test_natural_image_0(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(