Add model linearization, and expanded weights spatial positions

* Optionally replace non-linear MaxPool2d layers with their linear AvgPool2d equivalents.

* Added info for how to visualize expanded weights spatial positions in expanded weights / weight vis tutorial.

Author: ProGamerGov

captum/optim/_utils/circuits.py

@@ -1,11 +1,11 @@
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 
 from captum.optim._param.image.transform import center_crop_shape
 from captum.optim._utils.models import collect_activations
-from captum.optim._utils.typing import ModelInputType, TransformSize
+from captum.optim._utils.typing import ModelInputType, PoolParam, TransformSize
 
 
 def get_expanded_weights(
@@ -56,8 +56,57 @@ def get_expanded_weights(
             retain_graph=True,
         )[0]
         A.append(x.squeeze(0))
-    exapnded_weights = torch.stack(A, 0)
+    expanded_weights = torch.stack(A, 0)
 
     if crop_shape is not None:
-        exapnded_weights = center_crop_shape(exapnded_weights, crop_shape)
-    return exapnded_weights
+        expanded_weights = center_crop_shape(expanded_weights, crop_shape)
+    return expanded_weights
+
+
+def max2avg_pool2d(model, value: Optional[Any] = float("-inf")) -> None:
+    """
+    Replace all non-linear MaxPool2d layers with their linear AvgPool2d equivalents.
+    This allows us to ignore non-linear 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 AvgPool2dInf(torch.nn.Module):
+        def __init__(
+            self,
+            kernel_size: PoolParam = 2,
+            stride: Optional[PoolParam] = 2,
+            padding: PoolParam = 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 = AvgPool2dInf(
+                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)

captum/optim/_utils/typing.py

@@ -39,3 +39,4 @@ def cleanup(self):
 TransformVal = Union[int, float, Tensor]
 TransformSize = Union[List[int], Tuple[int], int]
 ModelInputType = Union[Tuple[Tensor], Tensor]
+PoolParam = Union[int, Tuple[int, ...]]

tests/optim/utils/test_circuits.py

@@ -2,10 +2,11 @@
 import unittest
 
 import torch
+import torch.nn.functional as F
 
 import captum.optim._utils.circuits as circuits
 from captum.optim._models.inception_v1 import googlenet
-from tests.helpers.basic import BaseTest
+from tests.helpers.basic import BaseTest, assertTensorAlmostEqual
 
 
 class TestGetExpandedWeights(BaseTest):
@@ -46,5 +47,24 @@ def test_get_expanded_weights_crop_two_int(self) -> None:
         self.assertEqual(list(output_tensor.shape), [480, 256, 5, 5])
 
 
+class TestMax2AvgPool2d(BaseTest):
+    def test_max2avg_pool2d(self) -> None:
+        model = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=0)
+        )
+
+        circuits.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)
+
+
 if __name__ == "__main__":
     unittest.main()