Merge pull request #496 from greentfrapp/optim-wip

Preliminary PR on optim-wip

Author: NarineK

captum/optim/__init__.py

@@ -1,27 +1,8 @@
-from typing import Dict, Optional, Union, Callable, Iterable
-from typing_extensions import Protocol
-
-import torch
-import torch.nn as nn
-
-ParametersForOptimizers = Iterable[Union[torch.Tensor, Dict[str, torch.tensor]]]
-
-
-class HasLoss(Protocol):
-    def loss(self) -> torch.Tensor:
-        ...
-
-
-class Parameterized(Protocol):
-    parameters: ParametersForOptimizers
-
-
-class Objective(Parameterized, HasLoss):
-    def cleanup(self):
-        pass
-
-
-ModuleOutputMapping = Dict[nn.Module, Optional[torch.Tensor]]
-
-StopCriteria = Callable[[int, Objective, torch.optim.Optimizer], bool]
-
+"""optim submodule."""
+
+import captum.optim._core.objectives as objectives  # noqa: F401
+import captum.optim._param.image.images as images  # noqa: F401
+import captum.optim._param.image.transform as transform  # noqa: F401
+import captum.optim._utils.typing as typing  # noqa: F401
+from captum.optim._core.objectives import InputOptimization  # noqa: F401
+from captum.optim._param.image.images import ImageTensor  # noqa: F401

captum/optim/techs/__init__.py renamed to captum/optim/_core/__init__.py

@@ -0,0 +1,258 @@
+"""captum.optim.objectives."""
+
+from contextlib import suppress
+from typing import Callable, Iterable, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from tqdm.auto import tqdm
+
+from captum.optim._core.output_hook import AbortForwardException, ModuleOutputsHook
+from captum.optim._param.image.images import InputParameterization, NaturalImage
+from captum.optim._param.image.transform import RandomAffine
+from captum.optim._utils.typing import (
+    LossFunction,
+    ModuleOutputMapping,
+    Objective,
+    Parameterized,
+    SingleTargetLossFunction,
+    StopCriteria,
+)
+
+
+class InputOptimization(Objective, Parameterized):
+    """
+    Core function that optimizes an input to maximize a target (aka objective).
+    This is similar to gradient-based methods for adversarial examples, such
+    as FGSM. The code for this was based on the implementation by the authors of Lucid.
+    For more details, see the following:
+        https://github.com/tensorflow/lucid
+        https://distill.pub/2017/feature-visualization/
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        input_param: Optional[InputParameterization],
+        transform: Optional[nn.Module],
+        target_modules: Iterable[nn.Module],
+        loss_function: LossFunction,
+    ):
+        r"""
+        Args:
+            model (nn.Module):  The reference to PyTorch model instance.
+            input_param (nn.Module, optional):  A module that generates an input,
+                        consumed by the model.
+            transform (nn.Module, optional):  A module that transforms or preprocesses
+                        the input before being passed to the model.
+            target_modules (iterable of nn.Module):  A list of targets, objectives that
+                        are used to compute the loss function.
+            loss_function (callable): The loss function to minimize during optimization
+                        optimization.
+        """
+        self.model = model
+        self.hooks = ModuleOutputsHook(target_modules)
+        self.input_param = input_param or NaturalImage((224, 224))
+        self.transform = transform or RandomAffine(scale=True, translate=True)
+        self.loss_function = loss_function
+
+    def loss(self) -> torch.Tensor:
+        r"""Compute loss value for current iteration.
+        Returns:
+            *tensor* representing **loss**:
+            - **loss** (*tensor*):
+                        Size of the tensor corresponds to the targets passed.
+        """
+        image = self.input_param()._t[None, ...]
+
+        if self.transform:
+            image = self.transform(image)
+
+        with suppress(AbortForwardException):
+            _unreachable = self.model(image)  # noqa: F841
+
+        # consume_outputs return the captured values and resets the hook's state
+        module_outputs = self.hooks.consume_outputs()
+        loss_value = self.loss_function(module_outputs)
+        return loss_value
+
+    def cleanup(self):
+        r"""Garbage collection, mainly removing hooks."""
+        self.hooks.remove_hooks()
+
+    # Targets are managed by ModuleOutputHooks; we mainly just want a convenient setter
+    @property
+    def targets(self):
+        return self.hooks.targets
+
+    @targets.setter
+    def targets(self, value):
+        self.hooks.remove_hooks()
+        self.hooks = ModuleOutputsHook(value)
+
+    def parameters(self):
+        return self.input_param.parameters()
+
+    def optimize(
+        self,
+        stop_criteria: Optional[StopCriteria] = None,
+        optimizer: Optional[optim.Optimizer] = None,
+    ):
+        r"""Optimize input based on loss function and objectives.
+        Args:
+            stop_criteria (StopCriteria, optional):  A function that is called
+                        every iteration and returns a bool that determines whether
+                        to stop the optimization.
+                        See captum.optim.typing.StopCriteria for details.
+            optimizer (Optimizer, optional):  An torch.optim.Optimizer used to
+                        optimize the input based on the loss function.
+        Returns:
+            *list* of *np.arrays* representing the **history**:
+            - **history** (*list*):
+                        A list of loss values per iteration.
+                        Length of the list corresponds to the number of iterations
+        """
+        stop_criteria = stop_criteria or n_steps(1024)
+        optimizer = optimizer or optim.Adam(self.parameters(), lr=0.025)
+        assert isinstance(optimizer, optim.Optimizer)
+
+        history = []
+        step = 0
+        while stop_criteria(step, self, history, optimizer):
+            optimizer.zero_grad()
+            loss_value = self.loss()
+            history.append(loss_value.cpu().detach().numpy())
+            (-1 * loss_value.mean()).backward()
+            optimizer.step()
+            step += 1
+
+        self.cleanup()
+        return history
+
+
+def n_steps(n: int) -> StopCriteria:
+    """StopCriteria generator that uses number of steps as a stop criteria.
+    Args:
+        n (int):  Number of steps to run optimization.
+    Returns:
+        *StopCriteria* callable
+    """
+    pbar = tqdm(total=n, unit="step")
+
+    def continue_while(step, obj, history, optim):
+        if len(history) > 0:
+            pbar.set_postfix({"Objective": f"{history[-1].mean():.1f}"}, refresh=False)
+        if step < n:
+            pbar.update()
+            return True
+        else:
+            pbar.close()
+            return False
+
+    return continue_while
+
+
+def channel_activation(target: nn.Module, channel_index: int) -> LossFunction:
+    def loss_function(targets_to_values: ModuleOutputMapping):
+        activations = targets_to_values[target]
+        assert activations is not None
+        # ensure channel_index is valid
+        assert channel_index < activations.shape[1]
+        # assume NCHW
+        # NOTE: not necessarily true e.g. for Linear layers
+        # assert len(activations.shape) == 4
+        return activations[:, channel_index, ...]
+
+    return loss_function
+
+
+def neuron_activation(
+    target: nn.Module, channel_index: int, x: int = None, y: int = None
+) -> LossFunction:
+    # ensure channel_index will be valid
+    assert channel_index < target.out_channels
+
+    def loss_function(targets_to_values: ModuleOutputMapping):
+        activations = targets_to_values[target]
+        assert activations is not None
+        assert len(activations.shape) == 4  # assume NCHW
+        _, _, H, W = activations.shape
+
+        if x is None:
+            _x = W // 2
+        else:
+            assert x < W
+            _x = x
+
+        if y is None:
+            _y = H // 2
+        else:
+            assert y < W
+            _y = y
+
+        return activations[:, channel_index, _x, _y]
+
+    return loss_function
+
+
+def single_target_objective(
+    target: nn.Module, loss_function: SingleTargetLossFunction
+) -> LossFunction:
+    def inner(targets_to_values: ModuleOutputMapping):
+        value = targets_to_values[target]
+        return loss_function(value)
+
+    return inner
+
+
+class SingleTargetObjective(Objective):
+    def __init__(
+        self,
+        model: nn.Module,
+        target: nn.Module,
+        loss_function: Callable[[torch.Tensor], torch.Tensor],
+    ):
+        super(SingleTargetObjective, self).__init__(model=model, targets=[target])
+        self.loss_function = loss_function
+
+    def loss(self, targets_to_values):
+        assert len(self.targets) == 1
+        target = self.targets[0]
+        target_value = targets_to_values[target]
+        loss_value = self.loss_function(target_value)
+        self.history.append(loss_value.sum().cpu().detach().numpy().squeeze().item())
+        return loss_value
+
+
+class MultiObjective(Objective):
+    def __init__(
+        self, objectives: List[Objective], weights: Optional[Iterable[float]] = None
+    ):
+        model = objectives[0].model
+        assert all(o.model == model for o in objectives)
+        targets = (target for objective in objectives for target in objective.targets)
+        super(MultiObjective, self).__init__(model=model, targets=targets)
+        self.objectives = objectives
+        self.weights = weights or len(objectives) * [1]
+
+    def loss(self, targets_to_values):
+        loss = (
+            objective.loss_function(targets_to_values) for objective in self.objectives
+        )
+        weighted = (loss * weight for weight in self.weights)
+        loss_value = sum(weighted)
+        self.history.append(loss_value.cpu().detach().numpy().squeeze().item())
+        return loss_value
+
+    @property
+    def histories(self) -> List[List[float]]:
+        return [objective.history for objective in self.objectives]
+
+
+# class ChannelObjective(SingleTargetObjective):
+#     def __init__(self, channel: int, *args, **kwargs):
+#         loss_function = lambda activation: activation[:, channel, :, :].mean()
+#         super(ChannelObjective, self).__init__(
+#             *args, loss_function=loss_function, **kwargs
+#         )

captum/optim/_core/objectives.py

@@ -1,10 +1,9 @@
+from typing import Iterable
 from warnings import warn
-from typing import Iterable, Dict, Optional
 
-import torch
 import torch.nn as nn
 
-from clarity.pytorch import ModuleOutputMapping
+# from clarity.pytorch import ModuleOutputMapping
 
 
 class AbortForwardException(Exception):
@@ -39,7 +38,8 @@ class ModuleReuseException(Exception):
 
 class ModuleOutputsHook:
     def __init__(self, target_modules: Iterable[nn.Module]):
-        self.outputs: ModuleOutputMapping = dict.fromkeys(target_modules, None)
+        # self.outputs: ModuleOutputMapping = dict.fromkeys(target_modules, None)
+        self.outputs = dict.fromkeys(target_modules, None)
         self.hooks = [
             module.register_forward_hook(self._forward_hook())
             for module in target_modules
@@ -59,17 +59,19 @@ def forward_hook(module, input, output):
                 self.outputs[module] = output
             else:
                 warn(
-                    f"Hook attached to {module} was called multiple times. As of 2019-11-22 please don't reuse nn.Modules in your models."
+                    f"Hook attached to {module} was called multiple times. "
+                    "As of 2019-11-22 please don't reuse nn.Modules in your models."
                 )
             if self.is_ready:
                 raise AbortForwardException("Forward hook called, all outputs saved.")
 
         return forward_hook
 
-    def consume_outputs(self) -> ModuleOutputMapping:
+    def consume_outputs(self):  # -> ModuleOutputMapping:
         if not self.is_ready:
             warn(
-                "Consume captured outputs, but not all requested target outputs have been captured yet!"
+                "Consume captured outputs, but not all requested target outputs "
+                "have been captured yet!"
             )
         outputs = self.outputs
         self._reset_outputs()
@@ -84,5 +86,5 @@ def remove_hooks(self):
             hook.remove()
 
     def __del__(self):
-        print(f"DEL HOOKS!: {list(self.outputs.keys())}")
+        # print(f"DEL HOOKS!: {list(self.outputs.keys())}")
         self.remove_hooks()

captum/optim/optim/output_hook.py renamed to captum/optim/_core/output_hook.py

@@ -0,0 +1 @@
+# from .inception_v1 import googlenet