Optim-wip: Add main Activation Atlas tutorial & functions (#782)

* Add main Activation Atlas classes & functions

* Add main Activation Atlas tutorial notebook

* Remove unsued import

* Changes based on feedback

* Underscore rotation class functions
* Fix & improve documentation.
* Revert weight heatmap variable name changes.
* Correct rotation transform bug.

* Revert ufmt change as it causes isort to fail

* Improve documentation of `AngledNeuronDirection` & atlas functions

* Improve atlas related documentation

* Move atlas.py to _utils/image & improve atlas docs

* Add sum_loss_list() function & correct target type hints

* Add `sum_loss_list()` to `optim/_core/loss.py` with tests.
* Replace `sum()` with `opt.loss.sum_loss_list()` in ActivationAtlas tutorial notebook.
* Correct `target` type hints for `Loss`, `BaseLoss`, and `CompositeLoss`. The correct type hint should be `Union[nn.Module, List[nn.Module]]` as that is what the code supports and uses.

* RandomRotation JIT support & other improvements

* Better way to handle torch version check with JIT

* The `torch.jit.is_scripting()` function isn't supported by earlier versions of PyTorch. So I've come up with a better solution that will still support earlier versions of PyTorch.
* Exposed `align_corners` parameter to `RandomRotation` class initialization.
* Improved `RandomRotation` class documentation.

* Use better scale type hint in RandomRotation init function

* Add torch.distributions support to RandomRotation

* Ludwig wanted this according to his original PR.
* Also fix mypy bug.

* Add assert & more tests for RandomRotation

* Added `torch.distribution` assert check and more extensive testing for the RandomRotation transform.

* Adding SkipTest to RandomRotation reflection

* Fix formatting error

* Changes to main atlas tutorial notebook based on feedback

* `extract_grid_vectors` -> `compute_avg_cell_samples`

* Improved documentation, and added better descriptions to the main atlas tutorial notebook.

* Remove unused type hint

* Improve whitening description in main activation atlas tutorial

* Spelling & grammar fixes

* Improve `calc_grid_indices` documentation

Author: ProGamerGov

captum/optim/__init__.py

@@ -6,6 +6,7 @@
 from captum.optim._param.image import images, transforms  # noqa: F401
 from captum.optim._param.image.images import ImageTensor  # noqa: F401
 from captum.optim._utils import circuits, reducer  # noqa: F401
+from captum.optim._utils.image import atlas  # noqa: F401
 from captum.optim._utils.image.common import (  # noqa: F401
     nchannels_to_rgb,
     save_tensor_as_image,
@@ -23,6 +24,7 @@
     "circuits",
     "models",
     "reducer",
+    "atlas",
     "nchannels_to_rgb",
     "save_tensor_as_image",
     "show",

captum/optim/_core/loss.py

@@ -1,7 +1,7 @@
 import functools
 import operator
 from abc import ABC, abstractmethod, abstractproperty
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -27,7 +27,7 @@ def __init__(self) -> None:
         super(Loss, self).__init__()
 
     @abstractproperty
-    def target(self) -> nn.Module:
+    def target(self) -> Union[nn.Module, List[nn.Module]]:
         pass
 
     @abstractmethod
@@ -140,7 +140,9 @@ def loss_fn(module: ModuleOutputMapping) -> torch.Tensor:
 
 class BaseLoss(Loss):
     def __init__(
-        self, target: nn.Module = [], batch_index: Optional[int] = None
+        self,
+        target: Union[nn.Module, List[nn.Module]] = [],
+        batch_index: Optional[int] = None,
     ) -> None:
         super(BaseLoss, self).__init__()
         self._target = target
@@ -150,7 +152,7 @@ def __init__(
             self._batch_index = (batch_index, batch_index + 1)
 
     @property
-    def target(self) -> nn.Module:
+    def target(self) -> Union[nn.Module, List[nn.Module]]:
         return self._target
 
     @property
@@ -160,7 +162,10 @@ def batch_index(self) -> Tuple:
 
 class CompositeLoss(BaseLoss):
     def __init__(
-        self, loss_fn: Callable, name: str = "", target: nn.Module = []
+        self,
+        loss_fn: Callable,
+        name: str = "",
+        target: Union[nn.Module, List[nn.Module]] = [],
     ) -> None:
         super(CompositeLoss, self).__init__(target)
         self.__name__ = name
@@ -499,6 +504,94 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         return _dot_cossim(self.direction, activations, cossim_pow=self.cossim_pow)
 
 
+@loss_wrapper
+class AngledNeuronDirection(BaseLoss):
+    """
+    Visualize a direction vector with an optional whitened activation vector to
+    unstretch the activation space. Compared to the traditional Direction objectives,
+    this objective places more emphasis on angle by optionally multiplying the dot
+    product by the cosine similarity.
+
+    When cossim_pow is equal to 0, this objective works as a euclidean
+    neuron objective. When cossim_pow is greater than 0, this objective works as a
+    cosine similarity objective. An additional whitened neuron direction vector
+    can optionally be supplied to improve visualization quality for some models.
+
+    More information on the algorithm this objective uses can be found here:
+    https://github.com/tensorflow/lucid/issues/116
+
+    This Lucid equivalents of this loss function can be found here:
+    https://github.com/tensorflow/lucid/blob/master/notebooks/
+    activation-atlas/activation-atlas-simple.ipynb
+    https://github.com/tensorflow/lucid/blob/master/notebooks/
+    activation-atlas/class-activation-atlas.ipynb
+
+    Like the Lucid equivalents, our implementation differs slightly from the
+    associated research paper.
+
+    Carter, et al., "Activation Atlas", Distill, 2019.
+    https://distill.pub/2019/activation-atlas/
+    """
+
+    def __init__(
+        self,
+        target: torch.nn.Module,
+        vec: torch.Tensor,
+        vec_whitened: Optional[torch.Tensor] = None,
+        cossim_pow: float = 4.0,
+        x: Optional[int] = None,
+        y: Optional[int] = None,
+        eps: float = 1.0e-4,
+        batch_index: Optional[int] = None,
+    ) -> None:
+        """
+        Args:
+            target (nn.Module): A target layer instance.
+            vec (torch.Tensor): A neuron direction vector to use.
+            vec_whitened (torch.Tensor, optional): A whitened neuron direction vector.
+            cossim_pow (float, optional): The desired cosine similarity power to use.
+            x (int, optional): Optionally provide a specific x position for the target
+                neuron.
+            y (int, optional): Optionally provide a specific y position for the target
+                neuron.
+            eps (float, optional): If cossim_pow is greater than zero, the desired
+                epsilon value to use for cosine similarity calculations.
+        """
+        BaseLoss.__init__(self, target, batch_index)
+        self.vec = vec.unsqueeze(0) if vec.dim() == 1 else vec
+        self.vec_whitened = vec_whitened
+        self.cossim_pow = cossim_pow
+        self.eps = eps
+        self.x = x
+        self.y = y
+        if self.vec_whitened is not None:
+            assert self.vec_whitened.dim() == 2
+        assert self.vec.dim() == 2
+
+    def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
+        activations = targets_to_values[self.target]
+        activations = activations[self.batch_index[0] : self.batch_index[1]]
+        assert activations.dim() == 4 or activations.dim() == 2
+        assert activations.shape[1] == self.vec.shape[1]
+        if activations.dim() == 4:
+            _x, _y = get_neuron_pos(
+                activations.size(2), activations.size(3), self.x, self.y
+            )
+            activations = activations[..., _x, _y]
+
+        vec = (
+            torch.matmul(self.vec, self.vec_whitened)[0]
+            if self.vec_whitened is not None
+            else self.vec
+        )
+        if self.cossim_pow == 0:
+            return activations * vec
+
+        dot = torch.mean(activations * vec)
+        cossims = dot / (self.eps + torch.sqrt(torch.sum(activations ** 2)))
+        return dot * torch.clamp(cossims, min=0.1) ** self.cossim_pow
+
+
 @loss_wrapper
 class TensorDirection(BaseLoss):
     """
@@ -590,6 +683,47 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         return activations
 
 
+def sum_loss_list(
+    loss_list: List,
+    to_scalar_fn: Callable[[torch.Tensor], torch.Tensor] = torch.mean,
+) -> CompositeLoss:
+    """
+    Summarize a large number of losses without recursion errors. By default using 300+
+    loss functions for a single optimization task will result in exceeding Python's
+    default maximum recursion depth limit. This function can be used to avoid the
+    recursion depth limit for tasks such as summarizing a large list of loss functions
+    with the built-in sum() function.
+
+    This function works similar to Lucid's optvis.objectives.Objective.sum() function.
+
+    Args:
+
+        loss_list (list): A list of loss function objectives.
+        to_scalar_fn (Callable): A function for converting loss function outputs to
+            scalar values, in order to prevent size mismatches.
+            Default: torch.mean
+
+    Returns:
+        loss_fn (CompositeLoss): A composite loss function containing all the loss
+            functions from `loss_list`.
+    """
+
+    def loss_fn(module: ModuleOutputMapping) -> torch.Tensor:
+        return sum([to_scalar_fn(loss(module)) for loss in loss_list])
+
+    name = "Sum(" + ", ".join([loss.__name__ for loss in loss_list]) + ")"
+    # Collect targets from losses
+    target = [
+        target
+        for targets in [
+            [loss.target] if not hasattr(loss.target, "__iter__") else loss.target
+            for loss in loss_list
+        ]
+        for target in targets
+    ]
+    return CompositeLoss(loss_fn, name=name, target=target)
+
+
 def default_loss_summarize(loss_value: torch.Tensor) -> torch.Tensor:
     """
     Helper function to summarize tensor outputs from loss functions.
@@ -617,7 +751,9 @@ def default_loss_summarize(loss_value: torch.Tensor) -> torch.Tensor:
     "Alignment",
     "Direction",
     "NeuronDirection",
+    "AngledNeuronDirection",
     "TensorDirection",
     "ActivationWeights",
+    "sum_loss_list",
     "default_loss_summarize",
 ]

captum/optim/_param/image/transforms.py

@@ -384,6 +384,152 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         return self.translate_tensor(input, insets)
 
 
+class RandomRotation(nn.Module):
+    """
+    Apply random rotation transforms on a NCHW tensor, using a sequence of degrees or
+    torch.distributions instance.
+    """
+
+    __constants__ = [
+        "degrees",
+        "mode",
+        "padding_mode",
+        "align_corners",
+        "_has_align_corners",
+        "_is_distribution",
+    ]
+
+    def __init__(
+        self,
+        degrees: NumSeqOrTensorType,
+        mode: str = "bilinear",
+        padding_mode: str = "zeros",
+        align_corners: bool = False,
+    ) -> None:
+        """
+        Args:
+
+            degrees (float, sequence, or torch.distribution): Tuple of degrees values
+                to randomly select from, or a torch.distributions instance.
+            mode (str, optional): Interpolation mode to use. See documentation of
+                F.grid_sample for more details. One of; "bilinear", "nearest", or
+                "bicubic".
+                Default: "bilinear"
+            padding_mode (str, optional): Padding mode for values that fall outside of
+                the grid. See documentation of F.grid_sample for more details. One of;
+                "zeros", "border", or "reflection".
+                Default: "zeros"
+            align_corners (bool, optional): Whether or not to align corners. See
+                documentation of F.affine_grid & F.grid_sample for more details.
+                Default: False
+        """
+        super().__init__()
+        if isinstance(degrees, torch.distributions.distribution.Distribution):
+            # Distributions are not supported by TorchScript / JIT yet
+            assert degrees.batch_shape == torch.Size([])
+            self.degrees_distribution = degrees
+            self._is_distribution = True
+            self.degrees = []
+        else:
+            assert hasattr(degrees, "__iter__")
+            if torch.is_tensor(degrees):
+                assert cast(torch.Tensor, degrees).dim() == 1
+                degrees = degrees.tolist()
+            assert len(degrees) > 0
+            self.degrees = [float(d) for d in degrees]
+            self._is_distribution = False
+
+        self.mode = mode
+        self.padding_mode = padding_mode
+        self.align_corners = align_corners
+        self._has_align_corners = torch.__version__ >= "1.3.0"
+
+    def _get_rot_mat(
+        self,
+        theta: float,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """
+        Create a rotation matrix tensor.
+
+        Args:
+
+            theta (float): The rotation value in degrees.
+
+        Returns:
+            **rot_mat** (torch.Tensor): A rotation matrix.
+        """
+        theta = theta * math.pi / 180.0
+        rot_mat = torch.tensor(
+            [
+                [math.cos(theta), -math.sin(theta), 0.0],
+                [math.sin(theta), math.cos(theta), 0.0],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        return rot_mat
+
+    def _rotate_tensor(self, x: torch.Tensor, theta: float) -> torch.Tensor:
+        """
+        Rotate an NCHW image tensor based on a specified degree value.
+
+        Args:
+
+            x (torch.Tensor): The NCHW image tensor to rotate.
+            theta (float): The amount to rotate the NCHW image, in degrees.
+
+        Returns:
+            **x** (torch.Tensor): A rotated NCHW image tensor.
+        """
+        rot_matrix = self._get_rot_mat(theta, x.device, x.dtype)[None, ...].repeat(
+            x.shape[0], 1, 1
+        )
+        if self._has_align_corners:
+            # Pass align_corners explicitly for torch >= 1.3.0
+            grid = F.affine_grid(rot_matrix, x.size(), align_corners=self.align_corners)
+            x = F.grid_sample(
+                x,
+                grid,
+                mode=self.mode,
+                padding_mode=self.padding_mode,
+                align_corners=self.align_corners,
+            )
+        else:
+            grid = F.affine_grid(rot_matrix, x.size())
+            x = F.grid_sample(x, grid, mode=self.mode, padding_mode=self.padding_mode)
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Randomly rotate an NCHW image tensor.
+
+        Args:
+
+            x (torch.Tensor): NCHW image tensor to randomly rotate.
+
+        Returns:
+            **x** (torch.Tensor): A randomly rotated NCHW image *tensor*.
+        """
+        assert x.dim() == 4
+        if self._is_distribution:
+            rotate_angle = float(self.degrees_distribution.sample().item())
+        else:
+            n = int(
+                torch.randint(
+                    low=0,
+                    high=len(self.degrees),
+                    size=[1],
+                    dtype=torch.int64,
+                    layout=torch.strided,
+                    device=x.device,
+                ).item()
+            )
+            rotate_angle = self.degrees[n]
+        return self._rotate_tensor(x, rotate_angle)
+
+
 class ScaleInputRange(nn.Module):
     """
     Multiplies the input by a specified multiplier for models with input ranges other
@@ -673,6 +819,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     "center_crop",
     "RandomScale",
     "RandomSpatialJitter",
+    "RandomRotation",
     "ScaleInputRange",
     "RGBToBGR",
     "GaussianSmoothing",