Optim wip - General fixes, SharedImage, & Weight Visualization

captum/optim/__init__.py

@@ -6,4 +6,6 @@
 from captum.optim._param.image import images  # noqa: F401
 from captum.optim._param.image import transform  # noqa: F401
 from captum.optim._param.image.images import ImageTensor  # noqa: F401
-from captum.optim._utils import models  # noqa: F401
+from captum.optim._utils import circuits, models, reducer  # noqa: F401
+from captum.optim._utils.image.common import nchannels_to_rgb  # noqa: F401
+from captum.optim._utils.image.common import weights_to_heatmap_2d  # noqa: F401

captum/optim/_core/loss.py

@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 
-from captum.optim._utils.images import get_neuron_pos
+from captum.optim._utils.image.common import get_neuron_pos
 from captum.optim._utils.typing import ModuleOutputMapping
 
 
@@ -66,18 +66,16 @@ def __init__(
         self.x = x
         self.y = y
 
-        # ensure channel_index will be valid
-        assert self.channel_index < self.target.out_channels
-
-    def _call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
         assert activations is not None
+        assert self.channel_index < activations.shape[1]
         assert len(activations.shape) == 4  # assume NCHW
         _x, _y = get_neuron_pos(
             activations.size(2), activations.size(3), self.x, self.y
         )
 
-        return activations[:, self.channel_index, _x, _y]
+        return activations[:, self.channel_index, _x : _x + 1, _y : _y + 1]
 
 
 class DeepDream(Loss):
@@ -98,7 +96,7 @@ class TotalVariation(Loss):
     https://arxiv.org/abs/1412.0035
     """
 
-    def _call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
         x_diff = activations[..., 1:, :] - activations[..., :-1, :]
         y_diff = activations[..., :, 1:] - activations[..., :, :-1]
@@ -115,7 +113,7 @@ def __init__(self, target: nn.Module, constant: float = 0.0) -> None:
         self.target = target
         self.constant = constant
 
-    def _call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
         return torch.abs(activations - self.constant).sum()
 
@@ -132,7 +130,7 @@ def __init__(
         self.constant = constant
         self.epsilon = epsilon
 
-    def _call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
         activations = (activations - self.constant).sum()
         return torch.sqrt(self.epsilon + activations)
@@ -145,7 +143,7 @@ class Diversity(Loss):
     https://distill.pub/2017/feature-visualization/#diversity
     """
 
-    def _call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
         return -sum(
             [
@@ -260,7 +258,7 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         return torch.cosine_similarity(self.direction, activations)
 
 
-class DirectionNeuron(Loss):
+class NeuronDirection(Loss):
     """
     Visualize a single (x, y) position for a direction vector.
     Carter, et al., "Activation Atlas", Distill, 2019.
@@ -271,16 +269,16 @@ def __init__(
         self,
         target: nn.Module,
         vec: torch.Tensor,
-        channel_index: int,
         x: Optional[int] = None,
         y: Optional[int] = None,
+        channel_index: Optional[int] = None,
     ) -> None:
         super(Loss, self).__init__()
         self.target = target
         self.direction = vec.reshape((1, -1, 1, 1))
-        self.channel_index = channel_index
         self.x = x
         self.y = y
+        self.channel_index = channel_index
 
     def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
@@ -290,8 +288,10 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         _x, _y = get_neuron_pos(
             activations.size(2), activations.size(3), self.x, self.y
         )
-        activations = activations[:, self.channel_index, _x, _y]
-        return torch.cosine_similarity(self.direction, activations[None, None, None])
+        activations = activations[:, :, _x : _x + 1, _y : _y + 1]
+        if self.channel_index is not None:
+            activations = activations[:, self.channel_index, ...][:, None, ...]
+        return torch.cosine_similarity(self.direction, activations)
 
 
 class TensorDirection(Loss):
@@ -361,7 +361,9 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
                 _x, _y = get_neuron_pos(
                     activations.size(2), activations.size(3), self.x, self.y
                 )
-                activations = activations[..., _x, _y].squeeze() * self.weights
+                activations = (
+                    activations[..., _x : _x + 1, _y : _y + 1].squeeze() * self.weights
+                )
             else:
                 activations = activations[
                     ..., self.y : self.y + self.wy, self.x : self.x + self.wx

captum/optim/_core/objectives.py

@@ -10,7 +10,10 @@
 try:
     from tqdm.auto import tqdm
 except (ImportError, AssertionError):
-    print("The tqdm package is required to use Captum's Optim library")
+    print(
+        "The tqdm package is required to use captum.optim's"
+        + " n_steps stop criteria with progress bar"
+    )
 
 from captum.optim._core.output_hook import AbortForwardException, ModuleOutputsHook
 from captum.optim._param.image.images import InputParameterization, NaturalImage
@@ -147,23 +150,32 @@ def optimize(
         return history
 
 
-def n_steps(n: int) -> StopCriteria:
+def n_steps(n: int, show_progress: bool = True) -> StopCriteria:
     """StopCriteria generator that uses number of steps as a stop criteria.
     Args:
         n (int):  Number of steps to run optimization.
+        show_progress (bool, optional):  Whether or not to show progress bar.
+            Default: True
     Returns:
         *StopCriteria* callable
     """
-    pbar = tqdm(total=n, unit="step")
+
+    if show_progress:
+        pbar = tqdm(total=n, unit=" step")
 
     def continue_while(step, obj, history, optim) -> bool:
         if len(history) > 0:
-            pbar.set_postfix({"Objective": f"{history[-1].mean():.1f}"}, refresh=False)
+            if show_progress:
+                pbar.set_postfix(
+                    {"Objective": f"{history[-1].mean():.1f}"}, refresh=False
+                )
         if step < n:
-            pbar.update()
+            if show_progress:
+                pbar.update()
             return True
         else:
-            pbar.close()
+            if show_progress:
+                pbar.close()
             return False
 
     return continue_while

captum/optim/_core/output_hook.py

@@ -1,9 +1,10 @@
-from typing import Iterable
+from contextlib import suppress
+from typing import Iterable, List, Union
 from warnings import warn
 
 import torch.nn as nn
 
-# from clarity.pytorch import ModuleOutputMapping
+from captum.optim._utils.typing import ModelInputType, ModuleOutputMapping
 
 
 class AbortForwardException(Exception):
@@ -88,3 +89,23 @@ def remove_hooks(self) -> None:
     def __del__(self) -> None:
         # print(f"DEL HOOKS!: {list(self.outputs.keys())}")
         self.remove_hooks()
+
+
+class ActivationFetcher:
+    """
+    Simple module for collecting activations from model targets.
+    """
+
+    def __init__(self, model, targets: Union[nn.Module, List[nn.Module]]) -> None:
+        super(ActivationFetcher, self).__init__()
+        self.model = model
+        self.layers = ModuleOutputsHook(targets)
+
+    def __call__(self, input_t: ModelInputType) -> ModuleOutputMapping:
+        try:
+            with suppress(AbortForwardException):
+                self.model(input_t)
+            activations = self.layers.consume_outputs()
+        finally:
+            self.layers.remove_hooks()
+        return activations

captum/optim/_models/inception_v1.py

@@ -67,7 +67,6 @@ def __init__(
             out_channels=64,
             kernel_size=(7, 7),
             stride=(2, 2),
-            padding=(3, 3),
             groups=1,
             bias=True,
         )
@@ -89,7 +88,6 @@ def __init__(
             out_channels=192,
             kernel_size=(3, 3),
             stride=(1, 1),
-            padding=(1, 1),
             groups=1,
             bias=True,
         )
@@ -132,6 +130,7 @@ def _transform_input(self, x: torch.Tensor) -> torch.Tensor:
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self._transform_input(x)
+        x = F.pad(x, (2, 3, 2, 3))
         x = self.conv1(x)
         x = self.conv1_relu(x)
         x = F.pad(x, (0, 1, 0, 1), value=float("-inf"))
@@ -140,6 +139,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         x = self.conv2(x)
         x = self.conv2_relu(x)
+        x = F.pad(x, (1, 1, 1, 1))
         x = self.conv3(x)
         x = self.conv3_relu(x)
         x = self.localresponsenorm2(x)
@@ -214,7 +214,6 @@ def __init__(
             out_channels=c3x3,
             kernel_size=(3, 3),
             stride=(1, 1),
-            padding=(1, 1),
             groups=1,
             bias=True,
         )
@@ -234,7 +233,6 @@ def __init__(
             out_channels=c5x5,
             kernel_size=(5, 5),
             stride=(1, 1),
-            padding=(2, 2),
             groups=1,
             bias=True,
         )
@@ -257,18 +255,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         c3x3 = self.conv_3x3_reduce(x)
         c3x3 = self.conv_3x3_reduce_relu(c3x3)
+        c3x3 = F.pad(c3x3, (1, 1, 1, 1))
         c3x3 = self.conv_3x3(c3x3)
         c3x3 = self.conv_3x3_relu(c3x3)
 
         c5x5 = self.conv_5x5_reduce(x)
         c5x5 = self.conv_5x5_reduce_relu(c5x5)
+        c5x5 = F.pad(c5x5, (2, 2, 2, 2))
         c5x5 = self.conv_5x5(c5x5)
         c5x5 = self.conv_5x5_relu(c5x5)
 
-        px = self.pool_proj(x)
-        px = self.pool_proj_relu(px)
-        px = F.pad(px, (1, 1, 1, 1), value=float("-inf"))
+        px = F.pad(x, (1, 1, 1, 1), value=float("-inf"))
         px = self.pool(px)
+        px = self.pool_proj(px)
+        px = self.pool_proj_relu(px)
         return torch.cat([c1x1, c3x3, c5x5, px], dim=1)