Add Activation Atlas tutorial & functions

Author: ProGamerGov

captum/optim/__init__.py

@@ -6,6 +6,7 @@
 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 circuits, models, reducer  # noqa: F401
+from captum.optim._utils import atlas, circuits, models, reducer  # noqa: F401
+from captum.optim._utils.image import dataset  # 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/_param/image/transform.py

@@ -267,6 +267,53 @@ 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.
+    Arguments:
+        degrees (float, sequence): Tuple of degrees to randomly select from.
+    """
+
+    def __init__(
+        self, degrees: Union[List[float], Tuple[float, ...], torch.Tensor]
+    ) -> None:
+        super().__init__()
+        assert hasattr(degrees, "__iter__")
+        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 * 3.141592653589793 / 180
+        rot_matrix = self.get_rot_mat(theta, x.device, x.dtype)[None, ...].repeat(
+            x.shape[0], 1, 1
+        )
+        grid = F.affine_grid(rot_matrix, x.size())
+        x = F.grid_sample(x, grid)
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.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

captum/optim/_utils/atlas.py

@@ -0,0 +1,114 @@
+from typing import List, Tuple
+
+import torch
+
+
+def grid_indices(
+    tensor: torch.Tensor,
+    size: Tuple[int, int] = (8, 8),
+    x_extent: Tuple[float, float] = (0.0, 1.0),
+    y_extent: Tuple[float, float] = (0.0, 1.0),
+) -> List[List[torch.Tensor]]:
+    """
+    Create grid cells of a specified size for an irregular grid.
+    """
+
+    assert tensor.dim() == 2 and tensor.size(1) == 2
+    x_coords = ((tensor[:, 0] - x_extent[0]) / (x_extent[1] - x_extent[0])) * size[1]
+    y_coords = ((tensor[:, 1] - y_extent[0]) / (y_extent[1] - y_extent[0])) * size[0]
+
+    x_list = []
+    for x in range(size[1]):
+        y_list = []
+        for y in range(size[0]):
+            in_bounds_x = torch.logical_and(x <= x_coords, x_coords <= x + 1)
+            in_bounds_y = torch.logical_and(y <= y_coords, y_coords <= y + 1)
+            in_bounds_indices = torch.where(
+                torch.logical_and(in_bounds_x, in_bounds_y)
+            )[0]
+            y_list.append(in_bounds_indices)
+        x_list.append(y_list)
+    return x_list
+
+
+def normalize_grid(
+    x: torch.Tensor,
+    min_percentile: float = 0.01,
+    max_percentile: float = 0.99,
+    relative_margin: float = 0.1,
+) -> torch.Tensor:
+    """
+    Remove outliers and rescale grid to [0,1].
+    """
+
+    assert x.dim() == 2 and x.size(1) == 2
+    mins = torch.quantile(x, min_percentile, dim=0)
+    maxs = torch.quantile(x, max_percentile, dim=0)
+
+    # add margins
+    mins = mins - relative_margin * (maxs - mins)
+    maxs = maxs + relative_margin * (maxs - mins)
+
+    clipped = torch.max(torch.min(x, maxs), mins)
+    clipped = clipped - clipped.min(0)[0]
+    return clipped / clipped.max(0)[0]
+
+
+def extract_grid_vectors(
+    grid: List[List[torch.Tensor]],
+    activations: torch.Tensor,
+    size: Tuple[int, int] = (8, 8),
+    min_density: int = 8,
+) -> Tuple[torch.Tensor, List[Tuple[int, int]]]:
+    """
+    Create direction vectors.
+    """
+
+    cell_coords = []
+    average_activations = []
+    for x in range(size[1]):
+        for y in range(size[0]):
+            indices = grid[x][y]
+            if len(indices) >= min_density:
+                average_activations.append(torch.mean(activations[indices], 0))
+                cell_coords.append((x, y))
+    return torch.stack(average_activations), cell_coords
+
+
+def create_atlas_vectors(
+    tensor: torch.Tensor,
+    activations: torch.Tensor,
+    size: Tuple[int, int] = (8, 8),
+    min_density: int = 8,
+    normalize: bool = True,
+) -> Tuple[torch.Tensor, List[Tuple[int, int]]]:
+    """
+    Create direction vectors by splitting an irregular grid into cells.
+    """
+
+    assert tensor.dim() == 2 and tensor.size(1) == 2
+    if normalize:
+        tensor = normalize_grid(tensor)
+    indices = grid_indices(tensor, size)
+    grid_vecs, vec_coords = extract_grid_vectors(
+        indices, activations, size, min_density
+    )
+    return grid_vecs, vec_coords
+
+
+def create_atlas(
+    cells: List[torch.Tensor],
+    coords: List[List[torch.Tensor]],
+    grid_size: Tuple[int, int] = (8, 8),
+) -> torch.Tensor:
+    cell_h, cell_w = cells[0].shape[2:]
+    canvas = torch.ones(1, 3, cell_h * grid_size[0], cell_w * grid_size[1])
+    for i, img in enumerate(cells):
+        y = int(coords[i][0])
+        x = int(coords[i][1])
+        canvas[
+            ...,
+            (grid_size[0] - x - 1) * cell_h : (grid_size[0] - x) * cell_h,
+            y * cell_w : (y + 1) * cell_w,
+        ] = img
+    return canvas

captum/optim/_utils/image/dataset.py

@@ -1,5 +1,10 @@
+from typing import Dict, List, Optional
+
 import torch
 
+from captum.optim._utils.models import collect_activations
+from captum.optim._utils.typing import ModuleOutputMapping
+
 
 def image_cov(tensor: torch.Tensor) -> torch.Tensor:
     """
@@ -51,3 +56,67 @@ def dataset_klt_matrix(
 
     cov_mtx = dataset_cov_matrix(loader)
     return cov_matrix_to_klt(cov_mtx, normalize)
+
+
+def capture_activation_samples(
+    loader: torch.utils.data.DataLoader,
+    model,
+    targets: List[torch.nn.Module],
+    target_names: List[str],
+    num_samples: Optional[int] = None,
+    input_device: torch.device = torch.device("cpu"),
+) -> ModuleOutputMapping:
+    """
+    Create a dict of randomly sampled activations for an image dataset.
+
+    Args:
+        loader (torch.utils.data.DataLoader): A torch.utils.data.DataLoader
+            instance for an image dataset.
+        model (nn.Module): A PyTorch model instance.
+        targets (list of nn.Module): A list of layers to sample activations
+            from.
+        target_names (list of str): A list of names to use for the layers
+            to targets in the output dict.
+        num_samples (int): How many samples to collect. Default is to collect
+            all samples.
+        input_device (torch.device): The device to use for model inputs.
+    Returns:
+        activation_dict (dict of tensor): A dictionary containing the sampled
+            dataset activations, with the target_names as the keys.
+    """
+
+    def random_sample(activations: torch.Tensor) -> torch.Tensor:
+        """
+        Randomly sample H & W dimensions of activations with 4 dimensions.
+        """
+
+        rnd_samples = []
+        for b in range(activations.size(0)):
+            if activations.dim() == 4:
+                h, w = activations.shape[2:]
+                y = torch.randint(low=1, high=h, size=[1])
+                x = torch.randint(low=1, high=w, size=[1])
+                activ = activations[b, :, y, x]
+            elif activations.dim() == 2:
+                activ = activations[b].unsqueeze(1)
+            rnd_samples.append(activ)
+        return torch.cat(rnd_samples, 1).permute(1, 0)
+
+    assert len(target_names) == len(targets)
+    activation_dict: Dict = {k: [] for k in dict.fromkeys(target_names).keys()}
+
+    sample_count = 0
+    with torch.no_grad():
+        for inputs, _ in loader:
+            inputs = inputs.to(input_device)
+            target_activ_dict = collect_activations(model, targets, inputs)
+            for t in target_activ_dict.keys():
+                target_activ_dict[t] = [random_sample(target_activ_dict[t])]
+            activation_dict = {
+                k: activation_dict[k] + target_activ_dict[k] for k in activation_dict
+            }
+            sample_count += inputs.size(0)
+            if num_samples is not None:
+                if sample_count > num_samples:
+                    return {k: torch.cat(activation_dict[k]) for k in activation_dict}
+    return {k: torch.cat(activation_dict[k]) for k in activation_dict}

tests/optim/utils/image/dataset.py

@@ -4,6 +4,7 @@
 import torch
 
 import captum.optim._utils.image.dataset as dataset_utils
+from captum.optim._models.inception_v1 import googlenet
 from tests.helpers.basic import (
     BaseTest,
     assertArraysAlmostEqual,
@@ -105,5 +106,28 @@ def create_tensor() -> torch.Tensor:
         assertTensorAlmostEqual(self, klt_transform, expected_mtx)
 
 
+class TestCaptureActivationSamples(BaseTest):
+    def test_capture_activation_samples(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping capture_activation_samples test due to"
+                + "insufficient Torch version."
+            )
+
+        num_tensors = 10
+        dataset_tensors = [torch.ones(3, 224, 224) for x in range(num_tensors)]
+        test_dataset = dataset_helpers.ImageTestDataset(dataset_tensors)
+        dataset_loader = torch.utils.data.DataLoader(
+            test_dataset, batch_size=10, num_workers=0, shuffle=False
+        )
+        model = googlenet(pretrained=True)
+        targets = [model.mixed4c]
+        target_names = ["mixed4c"]
+        activation_dict = dataset_utils.capture_activation_samples(
+            dataset_loader, model, targets, target_names
+        )
+        self.assertEqual(list(activation_dict["mixed4c"].shape), [num_tensors, 512])
+
+
 if __name__ == "__main__":
     unittest.main()