Add JIT support for SharedImage._interpolate_tensor

* Added JIT support for SharedImage's interpolation operations.
* Unfortunately, JIT support required me to separate SharedImage's bilinear and trilinear resizing into separate functions as Union's of tuples are currently broken. Union support was also a newer addition, so now SharedImage can support older PyTorch versions as well.

Author: ProGamerGov

captum/optim/_param/image/images.py

@@ -454,7 +454,12 @@ class SharedImage(AugmentedImageParameterization):
     https://distill.pub/2018/differentiable-parameterizations/
     """
 
-    __constants__ = ["offset", "_supports_is_scripting"]
+    __constants__ = [
+        "offset",
+        "_supports_is_scripting",
+        "_has_align_corners",
+        "_has_recompute_scale_factor",
+    ]
 
     def __init__(
         self,
@@ -491,6 +496,8 @@ def __init__(
 
         # Check & store whether or not we can use torch.jit.is_scripting()
         self._supports_is_scripting = torch.__version__ >= "1.6.0"
+        self._has_align_corners = torch.__version__ >= "1.3.0"
+        self._has_recompute_scale_factor = torch.__version__ >= "1.6.0"
 
     def _get_offset(self, offset: Union[int, Tuple[int]], n: int) -> List[List[int]]:
         """
@@ -551,7 +558,75 @@ def _apply_offset(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
             A.append(x)
         return A
 
-    @torch.jit.ignore
+    def _interpolate_bilinear(
+        self,
+        x: torch.Tensor,
+        size: Tuple[int, int],
+    ) -> torch.Tensor:
+        """
+        Perform interpolation without any warnings.
+
+        Args:
+
+            x (torch.Tensor): The NCHW tensor to resize.
+            size (tuple of int): The desired output size to resize the input
+                to, with a format of: [height, width].
+
+        Returns:
+            x (torch.Tensor): A resized NCHW tensor.
+        """
+        assert x.dim() == 4
+        assert len(size) == 2
+
+        if self._has_align_corners:
+            if self._has_recompute_scale_factor:
+                x = F.interpolate(
+                    x,
+                    size=size,
+                    mode="bilinear",
+                    align_corners=False,
+                    recompute_scale_factor=False,
+                )
+            else:
+                x = F.interpolate(x, size=size, mode="bilinear", align_corners=False)
+        else:
+            x = F.interpolate(x, size=size, mode="bilinear")
+        return x
+
+    def _interpolate_trilinear(
+        self,
+        x: torch.Tensor,
+        size: Tuple[int, int, int],
+    ) -> torch.Tensor:
+        """
+        Perform interpolation without any warnings.
+
+        Args:
+
+            x (torch.Tensor): The NCHW tensor to resize.
+            size (tuple of int): The desired output size to resize the input
+                to, with a format of: [channels, height, width].
+
+        Returns:
+            x (torch.Tensor): A resized NCHW tensor.
+        """
+        x = x.unsqueeze(0)
+        assert x.dim() == 5
+        if self._has_align_corners:
+            if self._has_recompute_scale_factor:
+                x = F.interpolate(
+                    x,
+                    size=size,
+                    mode="trilinear",
+                    align_corners=False,
+                    recompute_scale_factor=False,
+                )
+            else:
+                x = F.interpolate(x, size=size, mode="trilinear", align_corners=False)
+        else:
+            x = F.interpolate(x, size=size, mode="trilinear")
+        return x.squeeze(0)
+
     def _interpolate_tensor(
         self, x: torch.Tensor, batch: int, channels: int, height: int, width: int
     ) -> torch.Tensor:
@@ -572,21 +647,14 @@ def _interpolate_tensor(
         """
 
         if x.size(1) == channels:
-            mode = "bilinear"
             size = (height, width)
+            x = self._interpolate_bilinear(x, size=size)
         else:
-            mode = "trilinear"
-            x = x.unsqueeze(0)
             size = (channels, height, width)
-        x = F.interpolate(x, size=size, mode=mode)
-        x = x.squeeze(0) if len(size) == 3 else x
+            x = self._interpolate_trilinear(x, size=size)
         if x.size(0) != batch:
             x = x.permute(1, 0, 2, 3)
-            x = F.interpolate(
-                x.unsqueeze(0),
-                size=(batch, x.size(2), x.size(3)),
-                mode="trilinear",
-            ).squeeze(0)
+            x = self._interpolate_trilinear(x, size=(batch, x.size(2), x.size(3)))
             x = x.permute(1, 0, 2, 3)
         return x
 

tests/optim/param/test_images.py

@@ -424,6 +424,74 @@ def test_subclass(self) -> None:
             issubclass(images.SharedImage, images.AugmentedImageParameterization)
         )
 
+    def test_sharedimage_interpolate_bilinear(self) -> None:
+        shared_shapes = (128 // 2, 128 // 2)
+        test_param = lambda: torch.ones(3, 3, 224, 224)  # noqa: E731
+        image_param = images.SharedImage(
+            shapes=shared_shapes, parameterization=test_param
+        )
+
+        size = (224, 128)
+        test_input = torch.randn(1, 3, 128, 128)
+
+        test_output = image_param._interpolate_bilinear(test_input.clone(), size=size)
+        expected_output = torch.nn.functional.interpolate(
+            test_input.clone(), size=size, mode="bilinear"
+        )
+        assertTensorAlmostEqual(self, test_output, expected_output, 0.0)
+
+        size = (128, 128)
+        test_input = torch.randn(1, 3, 224, 224)
+
+        test_output = image_param._interpolate_bilinear(test_input.clone(), size=size)
+        expected_output = torch.nn.functional.interpolate(
+            test_input.clone(), size=size, mode="bilinear"
+        )
+        assertTensorAlmostEqual(self, test_output, expected_output, 0.0)
+
+    def test_sharedimage_interpolate_trilinear(self) -> None:
+        shared_shapes = (128 // 2, 128 // 2)
+        test_param = lambda: torch.ones(3, 3, 224, 224)  # noqa: E731
+        image_param = images.SharedImage(
+            shapes=shared_shapes, parameterization=test_param
+        )
+
+        size = (3, 224, 128)
+        test_input = torch.randn(1, 1, 128, 128)
+
+        test_output = image_param._interpolate_trilinear(test_input.clone(), size=size)
+        expected_output = torch.nn.functional.interpolate(
+            test_input.clone().unsqueeze(0), size=size, mode="trilinear"
+        ).squeeze(0)
+        assertTensorAlmostEqual(self, test_output, expected_output, 0.0)
+
+        size = (2, 128, 128)
+        test_input = torch.randn(1, 4, 224, 224)
+
+        test_output = image_param._interpolate_trilinear(test_input.clone(), size=size)
+        expected_output = torch.nn.functional.interpolate(
+            test_input.clone().unsqueeze(0), size=size, mode="trilinear"
+        ).squeeze(0)
+        assertTensorAlmostEqual(self, test_output, expected_output, 0.0)
+
+    def test_torch_version_check(self) -> None:
+        shared_shapes = (128 // 2, 128 // 2)
+        test_param = lambda: torch.ones(3, 3, 224, 224)  # noqa: E731
+        image_param = images.SharedImage(
+            shapes=shared_shapes, parameterization=test_param
+        )
+
+        has_align_corners = torch.__version__ >= "1.3.0"
+        self.assertEqual(image_param._has_align_corners, has_align_corners)
+
+        has_recompute_scale_factor = torch.__version__ >= "1.6.0"
+        self.assertEqual(
+            image_param._has_recompute_scale_factor, has_recompute_scale_factor
+        )
+
+        supports_is_scripting = torch.__version__ >= "1.6.0"
+        self.assertEqual(image_param._supports_is_scripting, supports_is_scripting)
+
     def test_sharedimage_get_offset_single_number(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(
@@ -772,6 +840,15 @@ def test_subclass(self) -> None:
             issubclass(images.StackImage, images.AugmentedImageParameterization)
         )
 
+    def test_stackimage_torch_version_check(self) -> None:
+        img_param_1 = images.SimpleTensorParameterization(torch.ones(1, 3, 4, 4))
+        img_param_2 = images.SimpleTensorParameterization(torch.ones(1, 3, 4, 4))
+        param_list = [img_param_1, img_param_2]
+        stack_param = images.StackImage(parameterizations=param_list)
+
+        supports_is_scripting = torch.__version__ >= "1.6.0"
+        self.assertEqual(stack_param._supports_is_scripting, supports_is_scripting)
+
     def test_stackimage_init(self) -> None:
         if torch.__version__ <= "1.2.0":
             raise unittest.SkipTest(