Optim-wip: Add model linearization, and expanded weights spatial positions

captum/optim/_param/image/transform.py

@@ -135,14 +135,22 @@ class CenterCrop(torch.nn.Module):
         pixels_from_edges (bool, optional): Whether to treat crop size
             values as the number of pixels from the tensor's edge, or an
             exact shape in the center.
+        offset_left (bool, optional): If the cropped away sides are not
+            equal in size, offset to the left. Default is set to False
+            for offseting to the right. This parameter is only valid when
+            pixels_from_edges is False.
     """
 
     def __init__(
-        self, size: IntSeqOrIntType = 0, pixels_from_edges: bool = False
+        self,
+        size: IntSeqOrIntType = 0,
+        pixels_from_edges: bool = False,
+        offset_left: bool = False,
     ) -> None:
         super(CenterCrop, self).__init__()
         self.crop_vals = size
         self.pixels_from_edges = pixels_from_edges
+        self.offset_left = offset_left
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """
@@ -153,11 +161,16 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             tensor (torch.Tensor): A center cropped tensor.
         """
 
-        return center_crop(input, self.crop_vals, self.pixels_from_edges)
+        return center_crop(
+            input, self.crop_vals, self.pixels_from_edges, self.offset_left
+        )
 
 
 def center_crop(
-    input: torch.Tensor, crop_vals: IntSeqOrIntType, pixels_from_edges: bool = False
+    input: torch.Tensor,
+    crop_vals: IntSeqOrIntType,
+    pixels_from_edges: bool = False,
+    offset_left: bool = False,
 ) -> torch.Tensor:
     """
     Center crop a specified amount from a tensor.
@@ -167,6 +180,10 @@ def center_crop(
         pixels_from_edges (bool, optional): Whether to treat crop size
             values as the number of pixels from the tensor's edge, or an
             exact shape in the center.
+        offset_left (bool, optional): If the cropped away sides are not
+            equal in size, offset to the left. Default is set to False
+            for offseting to the right. This parameter is only valid when
+            pixels_from_edges is False.
     Returns:
         *tensor*:  A center cropped tensor.
     """
@@ -188,8 +205,12 @@ def center_crop(
         sw, sh = w // 2 - (w_crop // 2), h // 2 - (h_crop // 2)
         x = input[..., sh : sh + h_crop, sw : sw + w_crop]
     else:
-        h_crop = h - int(round((h - crop_vals[0]) / 2.0))
-        w_crop = w - int(round((w - crop_vals[1]) / 2.0))
+        h_crop = h - int(math.ceil((h - crop_vals[0]) / 2.0))
+        w_crop = w - int(math.ceil((w - crop_vals[1]) / 2.0))
+        if h % 2 == 0 and crop_vals[0] % 2 != 0 or h % 2 != 0 and crop_vals[0] % 2 == 0:
+            h_crop = h_crop + 1 if offset_left else h_crop
+        if w % 2 == 0 and crop_vals[1] % 2 != 0 or w % 2 != 0 and crop_vals[1] % 2 == 0:
+            w_crop = w_crop + 1 if offset_left else w_crop
         x = input[..., h_crop - crop_vals[0] : h_crop, w_crop - crop_vals[1] : w_crop]
     return x
 

captum/optim/_utils/models.py

@@ -1,5 +1,5 @@
 import math
-from typing import List, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -70,8 +70,12 @@ class ReluLayer(nn.Module):
     Basic Hookable & Replaceable ReLU layer.
     """
 
+    def __init__(self, inplace: bool = True) -> None:
+        super(ReluLayer, self).__init__()
+        self.inplace = inplace
+
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        return F.relu(input, inplace=True)
+        return F.relu(input, inplace=self.inplace)
 
 
 def replace_layers(model, old_layer=ReluLayer, new_layer=RedirectedReluLayer) -> None:
@@ -168,3 +172,70 @@ def collect_activations(
     catch_activ = ActivationFetcher(model, targets)
     activ_out = catch_activ(model_input)
     return activ_out
+
+
+def max2avg_pool2d(model, value: Optional[Any] = float("-inf")) -> None:
+    """
+    Replace all nonlinear MaxPool2d layers with their linear AvgPool2d equivalents.
+    This allows us to ignore nonlinear values when calculating expanded weights.
+
+    Args:
+        model (nn.Module): A PyTorch model instance.
+        value (Any): Used to return any padding that's meant to be ignored by
+            pooling layers back to zero.
+    """
+
+    class AvgPool2dLayer(torch.nn.Module):
+        def __init__(
+            self,
+            kernel_size: Union[int, Tuple[int, ...]] = 2,
+            stride: Optional[Union[int, Tuple[int, ...]]] = 2,
+            padding: Union[int, Tuple[int, ...]] = 0,
+            ceil_mode: bool = False,
+            value: Optional[Any] = None,
+        ) -> None:
+            super().__init__()
+            self.avgpool = torch.nn.AvgPool2d(
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                ceil_mode=ceil_mode,
+            )
+            self.value = value
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            x = self.avgpool(x)
+            if self.value is not None:
+                x[x == self.value] = 0.0
+            return x
+
+    for name, child in model._modules.items():
+        if isinstance(child, torch.nn.MaxPool2d):
+            new_layer = AvgPool2dLayer(
+                kernel_size=child.kernel_size,
+                stride=child.stride,
+                padding=child.padding,
+                ceil_mode=child.ceil_mode,
+                value=value,
+            )
+            setattr(model, name, new_layer)
+        elif child is not None:
+            max2avg_pool2d(child)
+
+
+def ignore_layer(model, layer) -> None:
+    """
+    Replace target layers with layers that do nothing.
+    This is useful for removing the nonlinear ReLU
+    layers when creating expanded weights.
+
+    Args:
+        model (nn.Module): A PyTorch model instance.
+        layer (nn.Module): A layer class type.
+    """
+
+    class IgnoreLayer(torch.nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x
+
+    replace_layers(model, layer, IgnoreLayer)

tests/optim/helpers/numpy_transforms.py

@@ -1,3 +1,4 @@
+import math
 from typing import List, Optional, Tuple, Union, cast
 
 import numpy as np
@@ -63,14 +64,22 @@ class CenterCrop:
         pixels_from_edges (bool, optional): Whether to treat crop size values
             as the number of pixels from the tensor's edge, or an exact shape
             in the center.
+        offset_left (bool, optional): If the cropped away sides are not
+            equal in size, offset to the left. Default is set to False
+            for offseting to the right. This parameter is only valid when
+            pixels_from_edges is False.
     """
 
     def __init__(
-        self, size: IntSeqOrIntType = 0, pixels_from_edges: bool = False
+        self,
+        size: IntSeqOrIntType = 0,
+        pixels_from_edges: bool = False,
+        offset_left: bool = False,
     ) -> None:
         super(CenterCrop, self).__init__()
         self.crop_vals = size
         self.pixels_from_edges = pixels_from_edges
+        self.offset_left = offset_left
 
     def forward(self, input: np.ndarray) -> np.ndarray:
         """
@@ -81,11 +90,16 @@ def forward(self, input: np.ndarray) -> np.ndarray:
             tensor (array): A center cropped tensor.
         """
 
-        return center_crop(input, self.crop_vals, self.pixels_from_edges)
+        return center_crop(
+            input, self.crop_vals, self.pixels_from_edges, self.offset_left
+        )
 
 
 def center_crop(
-    input: np.ndarray, crop_vals: IntSeqOrIntType, pixels_from_edges: bool = False
+    input: np.ndarray,
+    crop_vals: IntSeqOrIntType,
+    pixels_from_edges: bool = False,
+    offset_left: bool = False,
 ) -> np.ndarray:
     """
     Center crop a specified amount from a array.
@@ -95,6 +109,10 @@ def center_crop(
         pixels_from_edges (bool, optional): Whether to treat crop size values
             as the number of pixels from the array's edge, or an exact shape
             in the center.
+        offset_left (bool, optional): If the cropped away sides are not
+            equal in size, offset to the left. Default is set to False
+            for offseting to the right. This parameter is only valid when
+            pixels_from_edges is False.
     Returns:
         *array*:  A center cropped array.
     """
@@ -116,8 +134,12 @@ def center_crop(
         sw, sh = w // 2 - (w_crop // 2), h // 2 - (h_crop // 2)
         x = input[..., sh : sh + h_crop, sw : sw + w_crop]
     else:
-        h_crop = h - int(round((h - crop_vals[0]) / 2.0))
-        w_crop = w - int(round((w - crop_vals[1]) / 2.0))
+        h_crop = h - int(math.ceil((h - crop_vals[0]) / 2.0))
+        w_crop = w - int(math.ceil((w - crop_vals[1]) / 2.0))
+        if h % 2 == 0 and crop_vals[0] % 2 != 0 or h % 2 != 0 and crop_vals[0] % 2 == 0:
+            h_crop = h_crop + 1 if offset_left else h_crop
+        if w % 2 == 0 and crop_vals[1] % 2 != 0 or w % 2 != 0 and crop_vals[1] % 2 == 0:
+            w_crop = w_crop + 1 if offset_left else w_crop
         x = input[..., h_crop - crop_vals[0] : h_crop, w_crop - crop_vals[1] : w_crop]
     return x
 

tests/optim/param/test_transforms.py

@@ -189,11 +189,11 @@ def test_center_crop_one_number_exact(self) -> None:
             [
                 torch.tensor(
                     [
-                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                        [0.0, 1.0, 1.0, 0.0, 1.0],
-                        [0.0, 1.0, 1.0, 0.0, 1.0],
-                        [0.0, 0.0, 0.0, 0.0, 1.0],
                         [1.0, 1.0, 1.0, 1.0, 1.0],
+                        [1.0, 0.0, 0.0, 0.0, 0.0],
+                        [1.0, 0.0, 1.0, 1.0, 0.0],
+                        [1.0, 0.0, 1.0, 1.0, 0.0],
+                        [1.0, 0.0, 0.0, 0.0, 0.0],
                     ]
                 )
             ]
@@ -224,6 +224,24 @@ def test_center_crop_two_numbers_exact(self) -> None:
         ).unsqueeze(0)
         assertTensorAlmostEqual(self, cropped_tensor, expected_tensor, 0)
 
+    def test_center_crop_offset_left_uneven_sides(self) -> None:
+        crop_mod = transform.CenterCrop(
+            [5, 5], pixels_from_edges=False, offset_left=True
+        )
+        x = torch.ones(1, 3, 5, 5)
+        px = F.pad(x, (5, 4, 5, 4), value=float("-inf"))
+        cropped_tensor = crop_mod(px)
+        assertTensorAlmostEqual(self, x, cropped_tensor)
+
+    def test_center_crop_offset_left_even_sides(self) -> None:
+        crop_mod = transform.CenterCrop(
+            [5, 5], pixels_from_edges=False, offset_left=True
+        )
+        x = torch.ones(1, 3, 5, 5)
+        px = F.pad(x, (5, 5, 5, 5), value=float("-inf"))
+        cropped_tensor = crop_mod(px)
+        assertTensorAlmostEqual(self, x, cropped_tensor)
+
 
 class TestCenterCropFunction(BaseTest):
     def test_center_crop_one_number(self) -> None:
@@ -286,11 +304,11 @@ def test_center_crop_one_number_exact(self) -> None:
             [
                 torch.tensor(
                     [
-                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                        [0.0, 1.0, 1.0, 0.0, 1.0],
-                        [0.0, 1.0, 1.0, 0.0, 1.0],
-                        [0.0, 0.0, 0.0, 0.0, 1.0],
                         [1.0, 1.0, 1.0, 1.0, 1.0],
+                        [1.0, 0.0, 0.0, 0.0, 0.0],
+                        [1.0, 0.0, 1.0, 1.0, 0.0],
+                        [1.0, 0.0, 1.0, 1.0, 0.0],
+                        [1.0, 0.0, 0.0, 0.0, 0.0],
                     ]
                 )
             ]
@@ -319,6 +337,22 @@ def test_center_crop_two_numbers_exact(self) -> None:
         ).unsqueeze(0)
         assertTensorAlmostEqual(self, cropped_tensor, expected_tensor)
 
+    def test_center_crop_offset_left_uneven_sides(self) -> None:
+        x = torch.ones(1, 3, 5, 5)
+        px = F.pad(x, (5, 4, 5, 4), value=float("-inf"))
+        cropped_tensor = transform.center_crop(
+            px, crop_vals=[5, 5], pixels_from_edges=False, offset_left=True
+        )
+        assertTensorAlmostEqual(self, x, cropped_tensor)
+
+    def test_center_crop_offset_left_even_sides(self) -> None:
+        x = torch.ones(1, 3, 5, 5)
+        px = F.pad(x, (5, 5, 5, 5), value=float("-inf"))
+        cropped_tensor = transform.center_crop(
+            px, crop_vals=[5, 5], pixels_from_edges=False, offset_left=True
+        )
+        assertTensorAlmostEqual(self, x, cropped_tensor)
+
 
 class TestBlendAlpha(BaseTest):
     def test_blend_alpha(self) -> None:

tests/optim/utils/test_circuits.py

@@ -5,6 +5,7 @@
 
 import captum.optim._utils.circuits as circuits
 from captum.optim._models.inception_v1 import googlenet
+from captum.optim._utils.models import RedirectedReluLayer, ignore_layer, max2avg_pool2d
 from tests.helpers.basic import BaseTest
 
 
@@ -45,6 +46,37 @@ def test_get_expanded_weights_crop_two_int(self) -> None:
         )
         self.assertEqual(list(output_tensor.shape), [480, 256, 5, 5])
 
+    def test_get_expanded_nonlinear_top_connections(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping get_expanded_weights nonlinear_top_connections test"
+                + " due to insufficient Torch version."
+            )
+
+        if torch.__version__ == "1.3.0":
+            norm_func = torch.norm
+        else:
+            norm_func = torch.linalg.norm
+        model = googlenet(pretrained=True)
+        max2avg_pool2d(model)
+        ignore_layer(model, RedirectedReluLayer)
+        output_tensor = circuits.get_expanded_weights(
+            model, model.pool3, model.mixed4a, 5
+        )
+        self.assertEqual(list(output_tensor.shape), [508, 480, 5, 5])
+
+        top_connected_neurons = torch.argsort(
+            torch.stack(
+                [
+                    -norm_func(output_tensor[i, 379, :, :])
+                    for i in range(output_tensor.shape[0])
+                ]
+            )
+        )[:10].tolist()
+
+        expected_list = [50, 437, 96, 398, 434, 423, 408, 436, 424, 168]
+        self.assertEqual(top_connected_neurons, expected_list)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/optim/utils/test_model_utils.py

@@ -317,5 +317,33 @@ def test_collect_activations(self) -> None:
         self.assertEqual(list(cast(torch.Tensor, m4d_activ).shape), [1, 528, 14, 14])
 
 
+class TestMax2AvgPool2d(BaseTest):
+    def test_max2avg_pool2d(self) -> None:
+        model = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=0)
+        )
+
+        model_utils.max2avg_pool2d(model)
+
+        test_tensor = torch.randn(128, 32, 16, 16)
+        test_tensor = F.pad(test_tensor, (0, 1, 0, 1), value=float("-inf"))
+        out_tensor = model(test_tensor)
+
+        avg_pool = torch.nn.AvgPool2d(kernel_size=3, stride=2, padding=0)
+        expected_tensor = avg_pool(test_tensor)
+        expected_tensor[expected_tensor == float("-inf")] = 0.0
+
+        assertTensorAlmostEqual(self, out_tensor, expected_tensor, 0)
+
+
+class TestIgnoreLayer(BaseTest):
+    def test_ignore_layer(self) -> None:
+        model = torch.nn.Sequential(torch.nn.ReLU())
+        x = torch.randn(1, 3, 4, 4)
+        model_utils.ignore_layer(model, torch.nn.ReLU)
+        output_tensor = model(x)
+        assertTensorAlmostEqual(self, x, output_tensor, 0)
+
+
 if __name__ == "__main__":
     unittest.main()