Optim-wip: Miscellaneous Fixes

captum/optim/__init__.py

@@ -7,6 +7,7 @@
 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 import dataset  # noqa: F401
 from captum.optim._utils.image.common import (  # noqa: F401
     hue_to_rgb,
     make_grid_image,
@@ -28,6 +29,7 @@
     "reducer",
     "make_grid_image",
     "atlas",
+    "dataset",
     "hue_to_rgb",
     "nchannels_to_rgb",
     "save_tensor_as_image",

captum/optim/_core/optimization.py

@@ -1,5 +1,3 @@
-"""captum.optim.optimization."""
-
 import warnings
 from typing import Callable, Iterable, Optional
 
@@ -31,10 +29,24 @@ 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/
+    as :class:`FGSM <captum.robust.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/
+
+    Alias: ``captum.optim.InputOptimization``
+
+    Example::
+
+        >>> model = opt.models.googlenet(pretrained=True)
+        >>> loss_fn = opt.loss.LayerActivation(model.mixed4c)
+        >>> image = opt.images.NaturalImage(size=(224, 224))
+        >>> transform = opt.transforms.TransformationRobustness()
+        >>>
+        >>> obj = opt.InputOptimization(model, loss_fn, image, transform)
+        >>> history = obj.optimize(opt.optimization.n_steps(512))
+        >>> image().show(figsize=(10, 10)) # Display results
     """
 
     def __init__(
@@ -47,13 +59,32 @@ def __init__(
         r"""
         Args:
 
-            model (nn.Module, optional):  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.
-            loss_function (callable): The loss function to minimize during optimization
-                        optimization.
+            model (nn.Module, optional): The reference to PyTorch model instance. Set
+                to ``None`` for no model instance.
+            loss_function (Callable): The :mod:`Loss <.loss>` objective instance to
+                minimize during optimization.
+            input_param (InputParameterization, optional): A module that generates an
+                input, consumed by the model. Example: An
+                :mod:`ImageParameterization <captum.optim.images>` instance.
+            transform (nn.Module, optional): A module that transforms or preprocesses
+                the input before being passed to the model. Set to
+                :class:`torch.nn.Identity` for no transforms.
+
+        Instance variables that be used in the :func:`InputOptimization.optimize`
+        function, custom optimization functions, and StopCriteria functions:
+
+        Attributes:
+
+            model (torch.nn.Module): The given model instance given when initializing
+                ``InputOptimization``. If ``model`` was set to ``None`` during
+                initialization, then an instance of :class:`torch.nn.Identity` will be
+                returned.
+            input_param (InputParameterization): The given input parameterization
+                instance given when initializing ``InputOptimization``.
+            loss_function (Loss): The composable :mod:`Loss <.loss>` instance given
+                when initializing ``InputOptimization``.
+            transform (torch.nn.Module): The given transform instance given when
+                initializing ``InputOptimization``.
         """
         self.model = model or nn.Identity()
         # Grab targets from loss_function
@@ -76,9 +107,9 @@ 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.
+            tensor representing **loss**:
+            - **loss** (torch.Tensor): Size of the tensor corresponds to the targets
+                passed.
         """
         input_t = self.input_param()
 
@@ -95,7 +126,9 @@ def loss(self) -> torch.Tensor:
         return loss_value
 
     def cleanup(self) -> None:
-        r"""Garbage collection, mainly removing hooks."""
+        r"""Garbage collection, mainly removing hooks.
+        This should only be run after optimize is finished running.
+        """
         self.hooks.remove_hooks()
 
     # Targets are managed by ModuleOutputHooks; we mainly just want a convenient setter
@@ -109,6 +142,11 @@ def targets(self, value: Iterable[nn.Module]) -> None:
         self.hooks = ModuleOutputsHook(value)
 
     def parameters(self) -> Iterable[nn.Parameter]:
+        """
+        Returns:
+            parameters (iterable of torch.nn.Parameter): An iterable of parameters in
+                the input parameterization.
+        """
         return self.input_param.parameters()
 
     def optimize(
@@ -122,18 +160,19 @@ def optimize(
 
         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.
+            stop_criteria (StopCriteria, optional): A function that is called
+                every iteration and returns a bool that determines whether to stop the
+                optimization.
+                Default: :func:`n_steps(512) <.n_steps>`
+            optimizer (torch.optim.Optimizer, optional): A ``torch.optim.Optimizer``
+                instance to use for optimizing the input based on the loss function.
+                Default: :class:`torch.optim.Adam`
             loss_summarize_fn (Callable, optional): The function to use for summarizing
                 tensor outputs from loss functions.
-                Default: default_loss_summarize
-            lr: (float, optional): If no optimizer is given, then lr is used as the
+                Default: :func:`.default_loss_summarize`
+            lr (float, optional): If no optimizer is given, then lr is used as the
                 learning rate for the Adam optimizer.
-                Default: 0.025
+                Default: ``0.025``
 
         Returns:
             history (torch.Tensor): A stack of loss values per iteration. The size
@@ -163,13 +202,18 @@ def optimize(
 def n_steps(n: int, show_progress: bool = True) -> StopCriteria:
     """StopCriteria generator that uses number of steps as a stop criteria.
 
+    Example::
+
+        >>> stop_criteria = opt.optimization.n_steps(512, True)
+
     Args:
-        n (int):  Number of steps to run optimization.
-        show_progress (bool, optional):  Whether or not to show progress bar.
-            Default: True
+
+        n (int): Number of steps to run optimization.
+        show_progress (bool, optional): Whether or not to show progress bar.
+            Default: ``True``
 
     Returns:
-        *StopCriteria* callable
+        StopCriteria (Callable): A stop criteria function.
     """
 
     if show_progress:

captum/optim/_param/image/images.py

@@ -21,6 +21,29 @@
 
 
 class ImageTensor(torch.Tensor):
+    r"""
+    A subclass of :class:`torch.Tensor` that provides functions for easy loading,
+    saving, and displaying image tensors.
+
+    Alias: ``captum.optim.ImageTensor``
+
+    Example using file path or URL::
+
+        >>> image_tensor = opt.images.ImageTensor.load(<path/to/image_file>)
+        >>> image_tensor.export(filename="image_tensor.jpg")  # Save image(s)
+        >>> image_tensor.show()  # Displays image(s) via Matplotlib
+
+    Example using ``torch.Tensor``::
+
+        >>> image_tensor = torch.randn(1, 3, 224, 224)
+        >>> image_tensor = opt.images.ImageTensor(image_tensor)
+
+    Example using ``np.ndarray``::
+
+        >>> image_tensor = np.random.rand(1, 3, 224, 224)
+        >>> image_tensor = opt.images.ImageTensor(image_tensor)
+    """
+
     @staticmethod
     def __new__(
         cls: Type["ImageTensor"],
@@ -32,10 +55,10 @@ def __new__(
         Args:
 
             x (list or np.ndarray or torch.Tensor): A list, NumPy array, or PyTorch
-                tensor to create an `ImageTensor` from.
+                tensor to create an ``ImageTensor`` from.
 
         Returns:
-           x (ImageTensor): An `ImageTensor` instance.
+           x (ImageTensor): An ``ImageTensor`` instance.
         """
         if isinstance(x, torch.Tensor) and x.is_cuda:
             x.show = MethodType(cls.show, x)
@@ -45,17 +68,18 @@ def __new__(
             return super().__new__(cls, x, *args, **kwargs)
 
     @classmethod
-    def open(cls, path: str, scale: float = 255.0, mode: str = "RGB") -> "ImageTensor":
+    def load(cls, path: str, scale: float = 255.0, mode: str = "RGB") -> "ImageTensor":
         """
-        Load an image file from a URL or local filepath directly into an `ImageTensor`.
+        Load an image file from a URL or local filepath directly into an
+        ``ImageTensor``.
 
         Args:
 
             path (str): A URL or filepath to an image.
             scale (float, optional): The image scale to use.
-                Default: 255.0
+                Default: ``255.0``
             mode (str, optional): The image loading mode / colorspace to use.
-                Default: "RGB"
+                Default: ``"RGB"``
 
         Returns:
            x (ImageTensor): An `ImageTensor` instance.
@@ -69,6 +93,11 @@ def open(cls, path: str, scale: float = 255.0, mode: str = "RGB") -> "ImageTenso
         img_np = np.array(img.convert(mode)).astype(np.float32)
         return cls(img_np.transpose(2, 0, 1) / scale)
 
+    @classmethod
+    def open(cls, path: str, scale: float = 255.0, mode: str = "RGB") -> "ImageTensor":
+        r"""Alias for :func:`load`."""
+        return cls.load(path=path, scale=scale, mode=mode)
+
     def __repr__(self) -> str:
         prefix = "ImageTensor("
         indent = len(prefix)
@@ -104,24 +133,27 @@ def show(
         pad_value: float = 0.0,
     ) -> None:
         """
-        Display an `ImageTensor`.
+        Display image(s) in the ``ImageTensor`` instance using
+        :func:`captum.optim.show`.
 
         Args:
 
-            figsize (Tuple[int, int], optional): height & width to use
-                for displaying the `ImageTensor` figure.
-            scale (float, optional): Value to multiply the `ImageTensor` by so that
+            figsize (tuple of int, optional): The height & width to use for displaying
+                the ``ImageTensor`` figure, in the format of: (height, width).
+                Default: ``None``
+            scale (float, optional): Value to multiply the ``ImageTensor`` by so that
                 it's value range is [0-255] for display.
-                Default: 255.0
+                Default: ``255.0``
             images_per_row (int, optional): The number of images per row to use for the
-                grid image. Default is set to None for no grid image creation.
-                Default: None
+                grid image. Default is set to ``None`` for no grid image creation.
+                Default: ``None``
             padding (int, optional): The amount of padding between images in the grid
-                images. This parameter only has an effect if `nrow` is not None.
-                Default: 2
+                images. This parameter only has an effect if ``images_per_row`` is not
+                ``None``.
+                Default: ``2``
             pad_value (float, optional): The value to use for the padding. This
-                parameter only has an effect if `nrow` is not None.
-                Default: 0.0
+                parameter only has an effect if ``images_per_row`` is not None.
+                Default: ``0.0``
         """
         show(
             self,
@@ -142,27 +174,29 @@ def export(
         pad_value: float = 0.0,
     ) -> None:
         """
-        Save an `ImageTensor` as an image file.
+        Save image(s) in the `ImageTensor` instance as an image file, using
+        :func:`captum.optim.save_tensor_as_image`.
 
         Args:
 
-            filename (str): The filename to use when saving the `ImageTensor` as an
+            filename (str): The filename to use when saving the ``ImageTensor`` as an
                 image file.
-            scale (float, optional): Value to multiply the `ImageTensor` by so that
+            scale (float, optional): Value to multiply the ``ImageTensor`` by so that
                 it's value range is [0-255] for saving.
-                Default: 255.0
+                Default: ``255.0``
             mode (str, optional): A PIL / Pillow supported colorspace. Default is
                 set to None for automatic RGB / RGBA detection and usage.
-                Default: None
+                Default: ``None``
             images_per_row (int, optional): The number of images per row to use for the
                 grid image. Default is set to None for no grid image creation.
-                Default: None
+                Default: ``None``
             padding (int, optional): The amount of padding between images in the grid
-                images. This parameter only has an effect if `nrow` is not None.
-                Default: 2
+                images. This parameter only has an effect if ``images_per_row`` is not
+                ``None``.
+                Default: ``2``
             pad_value (float, optional): The value to use for the padding. This
-                parameter only has an effect if `nrow` is not None.
-                Default: 0.0
+                parameter only has an effect if ``images_per_row`` is not ``None``.
+                Default: ``0.0``
         """
         save_tensor_as_image(
             self,

captum/optim/_utils/image/atlas.py

@@ -78,6 +78,7 @@ def calc_grid_indices(
     ]
 
     Args:
+
         xy_grid (torch.tensor): The xy coordinate grid activation samples, with a shape
             of: [n_points, 2].
         grid_size (Tuple[int, int]): The grid_size of grid cells to use. The grid_size
@@ -86,6 +87,7 @@ def calc_grid_indices(
             Default: (0.0, 1.0)
         y_extent (Tuple[float, float], optional): The y axis range to use.
             Default: (0.0, 1.0)
+
     Returns:
         indices (list of list of torch.Tensors): List of lists of grid indices
             stored inside tensors to use. Each 1D tensor of indices has a size of: