Optim-wip: Add Activation Atlas tutorial & functions

captum/optim/__init__.py

@@ -5,7 +5,7 @@
 from captum.optim._core.optimization import InputOptimization  # noqa: F401
 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 import atlas, circuits, reducer  # noqa: F401
 from captum.optim._utils.image.common import (  # noqa: F401
     nchannels_to_rgb,
     save_tensor_as_image,

captum/optim/_core/loss.py

@@ -448,14 +448,15 @@ def __init__(
         batch_index: Optional[int] = None,
     ) -> None:
         BaseLoss.__init__(self, target, batch_index)
-        self.direction = vec.reshape((1, -1, 1, 1))
+        assert vec.dim() == 2 or vec.dim() == 4
+        self.vec = vec.reshape((vec.size(0), -1, 1, 1)) if vec.dim() == 2 else vec
         self.cossim_pow = cossim_pow
 
     def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
-        assert activations.size(1) == self.direction.size(1)
+        assert activations.size(1) == self.vec.size(1)
         activations = activations[self.batch_index[0] : self.batch_index[1]]
-        return _dot_cossim(self.direction, activations, cossim_pow=self.cossim_pow)
+        return _dot_cossim(self.vec, activations, cossim_pow=self.cossim_pow)
 
 
 @loss_wrapper
@@ -477,7 +478,8 @@ def __init__(
         batch_index: Optional[int] = None,
     ) -> None:
         BaseLoss.__init__(self, target, batch_index)
-        self.direction = vec.reshape((1, -1, 1, 1))
+        assert vec.dim() == 2 or vec.dim() == 4
+        self.vec = vec.reshape((vec.size(0), -1, 1, 1)) if vec.dim() == 2 else vec
         self.x = x
         self.y = y
         self.channel_index = channel_index
@@ -496,7 +498,81 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         ]
         if self.channel_index is not None:
             activations = activations[:, self.channel_index, ...][:, None, ...]
-        return _dot_cossim(self.direction, activations, cossim_pow=self.cossim_pow)
+        return _dot_cossim(self.vec, 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.
+
+    Carter, et al., "Activation Atlas", Distill, 2019.
+    https://distill.pub/2019/activation-atlas/
+    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.
+    """
+
+    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:
+        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
@@ -515,16 +591,17 @@ def __init__(
         batch_index: Optional[int] = None,
     ) -> None:
         BaseLoss.__init__(self, target, batch_index)
-        self.direction = vec
+        assert vec.dim() == 4
+        self.vec = vec
         self.cossim_pow = cossim_pow
 
     def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
         activations = targets_to_values[self.target]
 
         assert activations.dim() == 4
 
-        H_direction, W_direction = self.direction.size(2), self.direction.size(3)
-        H_activ, W_activ = activations.size(2), activations.size(3)
+        H_direction, W_direction = self.vec.shape[2:]
+        H_activ, W_activ = activations.shape[2:]
 
         H = (H_activ - H_direction) // 2
         W = (W_activ - W_direction) // 2
@@ -535,7 +612,7 @@ def __call__(self, targets_to_values: ModuleOutputMapping) -> torch.Tensor:
             H : H + H_direction,
             W : W + W_direction,
         ]
-        return _dot_cossim(self.direction, activations, cossim_pow=self.cossim_pow)
+        return _dot_cossim(self.vec, activations, cossim_pow=self.cossim_pow)
 
 
 @loss_wrapper
@@ -617,6 +694,7 @@ def default_loss_summarize(loss_value: torch.Tensor) -> torch.Tensor:
     "Alignment",
     "Direction",
     "NeuronDirection",
+    "AngledNeuronDirection",
     "TensorDirection",
     "ActivationWeights",
     "default_loss_summarize",

captum/optim/_param/image/transforms.py

@@ -302,6 +302,10 @@ def __init__(self, scale: NumSeqOrTensorType) -> None:
             scale (float, sequence): Tuple of rescaling values to randomly select from.
         """
         super().__init__()
+        assert hasattr(scale, "__iter__")
+        if torch.is_tensor(scale):
+            assert cast(torch.Tensor, scale).dim() == 1
+        assert len(scale) > 0
         self.scale = scale
 
     def get_scale_mat(
@@ -384,6 +388,75 @@ 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.
+    """
+
+    def __init__(
+        self, degrees: Union[List[float], Tuple[float, ...], torch.Tensor]
+    ) -> None:
+        """
+        Args:
+
+            degrees (float, sequence): Tuple, List, or Tensor of degrees to randomly
+                select from.
+        """
+        super().__init__()
+        assert hasattr(degrees, "__iter__")
+        if torch.is_tensor(degrees):
+            assert cast(torch.Tensor, degrees).dim() == 1
+        assert len(degrees) > 0
+        self.degrees = degrees
+
+    def get_rot_mat(
+        self,
+        theta: Union[int, float, torch.Tensor],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        theta = torch.tensor(theta, device=device, dtype=dtype)
+        rot_mat = torch.tensor(
+            [
+                [torch.cos(theta), -torch.sin(theta), 0],
+                [torch.sin(theta), torch.cos(theta), 0],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        return rot_mat
+
+    def rotate_tensor(
+        self, x: torch.Tensor, theta: Union[int, float, torch.Tensor]
+    ) -> torch.Tensor:
+        theta = theta * math.pi / 180
+        rot_matrix = self.get_rot_mat(theta, x.device, x.dtype)[None, ...].repeat(
+            x.shape[0], 1, 1
+        )
+        if torch.__version__ >= "1.3.0":
+            # Pass align_corners explicitly for torch >= 1.3.0
+            grid = F.affine_grid(rot_matrix, x.size(), align_corners=False)
+            x = F.grid_sample(x, grid, align_corners=False)
+        else:
+            grid = F.affine_grid(rot_matrix, x.size())
+            x = F.grid_sample(x, grid)
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Randomly rotate an input tensor.
+
+        Args:
+
+            input (torch.Tensor): Input to randomly rotate.
+
+        Returns:
+            **tensor** (torch.Tensor): A randomly rotated *tensor*.
+        """
+        rotate_angle = _rand_select(self.degrees)
+        return self.rotate_tensor(x, rotate_angle)
+
+
 class ScaleInputRange(nn.Module):
     """
     Multiplies the input by a specified multiplier for models with input ranges other