More doc improvements

Author: ProGamerGov

captum/optim/_param/image/transforms.py

@@ -20,7 +20,7 @@ def __init__(self, background: Optional[torch.Tensor] = None) -> None:
         """
         Args:
 
-            background (tensor, optional):  An NCHW image tensor to be used as the
+            background (tensor, optional): An NCHW image tensor to be used as the
                 Alpha channel's background.
                 Default: ``None``
         """
@@ -36,7 +36,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): RGBA image tensor to blend into an RGB image tensor.
 
         Returns:
-            **blended** (torch.Tensor): RGB image tensor.
+            blended (torch.Tensor): RGB image tensor.
         """
         assert x.dim() == 4
         assert x.size(1) == 4
@@ -60,7 +60,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): RGBA image tensor.
 
         Returns:
-            **rgb** (torch.Tensor): RGB image tensor without the alpha channel.
+            rgb (torch.Tensor): RGB image tensor without the alpha channel.
         """
         assert x.dim() == 4
         assert x.size(1) == 4
@@ -101,7 +101,7 @@ def klt_transform() -> torch.Tensor:
         Karhunen-Loève transform (KLT) measured on ImageNet
 
         Returns:
-            **transform** (torch.Tensor): A Karhunen-Loève transform (KLT) measured on
+            transform (torch.Tensor): A Karhunen-Loève transform (KLT) measured on
                 the ImageNet dataset.
         """
         # Handle older versions of PyTorch
@@ -120,7 +120,7 @@ def klt_transform() -> torch.Tensor:
     def i1i2i3_transform() -> torch.Tensor:
         """
         Returns:
-            **transform** (torch.Tensor): An approximation of natural colors transform
+            transform (torch.Tensor): An approximation of natural colors transform
                 (i1i2i3).
         """
         i1i2i3_matrix = [
@@ -134,7 +134,7 @@ def __init__(self, transform: Union[str, torch.Tensor] = "klt") -> None:
         """
         Args:
 
-            transform (str or tensor):  Either a string for one of the precalculated
+            transform (str or tensor): Either a string for one of the precalculated
                 transform matrices, or a 3x3 matrix for the 3 RGB channels of input
                 tensors.
         """
@@ -352,7 +352,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             input (torch.Tensor): Input to center crop.
 
         Returns:
-            **tensor** (torch.Tensor): A center cropped *tensor*.
+            tensor (torch.Tensor): A center cropped NCHW tensor.
         """
 
         return center_crop(
@@ -402,7 +402,7 @@ def center_crop(
             Default: ``0.0``
 
     Returns:
-        **tensor** (torch.Tensor):  A center cropped *tensor*.
+        tensor (torch.Tensor): A center cropped NCHW tensor.
     """
 
     assert input.dim() == 3 or input.dim() == 4
@@ -537,7 +537,7 @@ def _scale_tensor(self, x: torch.Tensor, scale: float) -> torch.Tensor:
             scale (float): The amount to scale the NCHW image by.
 
         Returns:
-            **x** (torch.Tensor): A scaled NCHW image tensor.
+            x (torch.Tensor): A scaled NCHW image tensor.
         """
         if self._has_antialias:
             x = F.interpolate(
@@ -567,7 +567,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): NCHW image tensor to randomly scale.
 
         Returns:
-            **x** (torch.Tensor): A randomly scaled NCHW image *tensor*.
+            x (torch.Tensor): A randomly scaled NCHW image tensor.
         """
         assert x.dim() == 4
         if self._is_distribution:
@@ -669,7 +669,7 @@ def _get_scale_mat(
             m (float): The scale value to use.
 
         Returns:
-            **scale_mat** (torch.Tensor): A scale matrix.
+            scale_mat (torch.Tensor): A scale matrix.
         """
         scale_mat = torch.tensor(
             [[m, 0.0, 0.0], [0.0, m, 0.0]], device=device, dtype=dtype
@@ -686,7 +686,7 @@ def _scale_tensor(self, x: torch.Tensor, scale: float) -> torch.Tensor:
             scale (float): The amount to scale the NCHW image by.
 
         Returns:
-            **x** (torch.Tensor): A scaled NCHW image tensor.
+            x (torch.Tensor): A scaled NCHW image tensor.
         """
         scale_matrix = self._get_scale_mat(scale, x.device, x.dtype)[None, ...].repeat(
             x.shape[0], 1, 1
@@ -710,7 +710,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): NCHW image tensor to randomly scale.
 
         Returns:
-            **x** (torch.Tensor): A randomly scaled NCHW image *tensor*.
+            x (torch.Tensor): A randomly scaled NCHW image tensor.
         """
         assert x.dim() == 4
         if self._is_distribution:
@@ -768,7 +768,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             input (torch.Tensor): Input to randomly translate.
 
         Returns:
-            **tensor** (torch.Tensor): A randomly translated *tensor*.
+            tensor (torch.Tensor): A randomly translated NCHW tensor.
         """
         insets = torch.randint(
             high=self.pad_range,
@@ -854,7 +854,7 @@ def _get_rot_mat(
             theta (float): The rotation value in degrees.
 
         Returns:
-            **rot_mat** (torch.Tensor): A rotation matrix.
+            rot_mat (torch.Tensor): A rotation matrix.
         """
         theta = theta * math.pi / 180.0
         rot_mat = torch.tensor(
@@ -877,7 +877,7 @@ def _rotate_tensor(self, x: torch.Tensor, theta: float) -> torch.Tensor:
             theta (float): The amount to rotate the NCHW image, in degrees.
 
         Returns:
-            **x** (torch.Tensor): A rotated NCHW image tensor.
+            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
@@ -901,7 +901,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): NCHW image tensor to randomly rotate.
 
         Returns:
-            **x** (torch.Tensor): A randomly rotated NCHW image *tensor*.
+            x (torch.Tensor): A randomly rotated NCHW image tensor.
         """
         assert x.dim() == 4
         if self._is_distribution:
@@ -933,7 +933,7 @@ def __init__(self, multiplier: float = 1.0) -> None:
         """
         Args:
 
-            multiplier (float, optional):  A float value used to scale the input.
+            multiplier (float, optional): A float value used to scale the input.
         """
         super().__init__()
         self.multiplier = multiplier
@@ -947,7 +947,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): Input to scale values of.
 
         Returns:
-            **tensor** (torch.Tensor): tensor with it's values scaled.
+            tensor (torch.Tensor): tensor with it's values scaled.
         """
         return x * self.multiplier
 
@@ -966,7 +966,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): RGB image tensor to convert to BGR.
 
         Returns:
-            **BGR tensor** (torch.Tensor): A BGR tensor.
+            BGR tensor (torch.Tensor): A BGR tensor.
         """
         assert x.dim() == 4
         assert x.size(1) == 3
@@ -1104,7 +1104,7 @@ def forward(
             x (torch.Tensor): Input to apply symmetric padding on.
 
         Returns:
-            **tensor** (torch.Tensor): Padded tensor.
+            tensor (torch.Tensor): Padded tensor.
         """
         ctx.padding = padding
         x_device = x.device
@@ -1127,7 +1127,7 @@ def backward(
             grad_output (torch.Tensor): Input to remove symmetric padding from.
 
         Returns:
-            **grad_input** (torch.Tensor): Unpadded tensor.
+            grad_input (torch.Tensor): Unpadded tensor.
         """
         grad_input = grad_output.clone()
         B, C, H, W = grad_input.size()
@@ -1166,7 +1166,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): Input to reduce channel dimensions on.
 
         Returns:
-            **3 channel RGB tensor** (torch.Tensor): RGB image tensor.
+            x (torch.Tensor): A 3 channel RGB image tensor.
         """
         assert x.dim() == 4
         return nchannels_to_rgb(x, self.warp)
@@ -1216,6 +1216,16 @@ def _center_crop(self, x: torch.Tensor) -> torch.Tensor:
         ]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Randomly crop an NCHW image tensor.
+
+        Args:
+
+            x (torch.Tensor): The NCHW image tensor to randomly crop.
+
+        Returns
+            x (torch.Tensor): The randomly cropped NCHW image tensor.
+        """
         assert x.dim() == 4
         hs = int(math.ceil((x.shape[2] - self.crop_size[0]) / 2.0))
         ws = int(math.ceil((x.shape[3] - self.crop_size[1]) / 2.0))

captum/optim/_utils/circuits.py

@@ -20,6 +20,24 @@ def extract_expanded_weights(
     literally adjacent in a neural network, or where the weights aren’t directly
     represented in a single weight tensor.
 
+    Example::
+
+        >>> # Load InceptionV1 model with nonlinear layers replaced by
+        >>> # their linear equivalents
+        >>> linear_model = opt.models.googlenet(
+        >>>     pretrained=True, use_linear_modules_only=True
+        >>> ).eval()
+        >>> # Extract weight interactions between target layers
+        >>> W_3a_3b = opt.circuits.extract_expanded_weights(
+        >>>     linear_model, linear_model.mixed3a, linear_model.mixed3b, 5
+        >>> )
+        >>> # Display results for channel 147 of mixed3a and channel 379 of
+        >>> # mixed3b, in human readable format
+        >>> W_3a_3b_hm = opt.weights_to_heatmap_2d(
+        >>>     W_3a_3b[379, 147, ...] / W_3a_3b[379, ...].max()
+        >>> )
+        >>> opt.show(W_3a_3b_hm)
+
     Voss, et al., "Visualizing Weights", Distill, 2021.
     See: https://distill.pub/2020/circuits/visualizing-weights/
 

captum/optim/_utils/image/atlas.py

@@ -146,13 +146,14 @@ def compute_avg_cell_samples(
             Default: ``8``
 
     Returns:
-        cell_vecs (torch.tensor): A tensor containing all the direction vectors that
-            were created, stacked along the batch dimension with a shape of:
-            [n_vecs, n_channels].
-        cell_coords (list of Tuple[int, int, int]): List of coordinates for grid
-            spatial positions of each direction vector, and the number of samples used
-            for the cell. The list for each cell is in the format of:
-            [x_coord, y_coord, number_of_samples_used].
+        cell_vecs_and_cell_coords: A 2 element tuple of: ``(cell_vecs, cell_coords)``.
+            - cell_vecs (torch.tensor): A tensor containing all the direction vectors
+                  that were created, stacked along the batch dimension with a shape of:
+                  [n_vecs, n_channels].
+            - cell_coords (list of Tuple[int, int, int]): List of coordinates for grid
+                  spatial positions of each direction vector, and the number of samples
+                  used for the cell. The list for each cell is in the format of:
+                  [x_coord, y_coord, number_of_samples_used].
     """
     assert raw_samples.dim() == 2
 
@@ -205,13 +206,14 @@ def create_atlas_vectors(
             Default: ``(0.0, 1.0)``
 
     Returns:
-        grid_vecs (torch.tensor): A tensor containing all the direction vectors that
-            were created, stacked along the batch dimension, with a shape of:
-            [n_vecs, n_channels].
-        cell_coords (list of Tuple[int, int, int]): List of coordinates for grid
-            spatial positions of each direction vector, and the number of samples used
-            for the cell. The list for each cell is in the format of:
-            [x_coord, y_coord, number_of_samples_used].
+        grid_vecs_and_cell_coords: A 2 element tuple of: ``(grid_vecs, cell_coords)``.
+            - grid_vecs (torch.tensor): A tensor containing all the direction vectors
+                  that were created, stacked along the batch dimension, with a shape
+                  of: [n_vecs, n_channels].
+            - cell_coords (list of Tuple[int, int, int]): List of coordinates for grid
+                  spatial positions of each direction vector, and the number of samples
+                  used for the cell. The list for each cell is in the format of:
+                  [x_coord, y_coord, number_of_samples_used].
     """
 
     assert xy_grid.dim() == 2 and xy_grid.size(1) == 2

captum/optim/_utils/image/common.py

@@ -348,7 +348,7 @@ def weights_to_heatmap_2d(
             Default: ``["0571b0", "92c5de", "f7f7f7", "f4a582", "ca0020"]``
 
     Returns:
-        color_tensor (torch.Tensor):  A weight heatmap.
+        color_tensor (torch.Tensor): A weight heatmap.
     """
 
     assert weight.dim() == 2

captum/optim/_utils/reducer.py

@@ -21,6 +21,14 @@ class ChannelReducer:
 
     See here for more information: https://distill.pub/2018/building-blocks/
 
+    Example::
+
+        >>> reducer = opt.reducer.ChannelReducer(2, "NMF")
+        >>> x = torch.randn(1, 8, 128, 128).abs()
+        >>> output = reducer.fit_transform(x)
+        >>> print(output.shape)
+        torch.Size([1, 2, 128, 128])
+
     Args:
 
         n_components (int, optional): The number of channels to reduce the target
@@ -30,7 +38,7 @@ class ChannelReducer:
             from sklearn, which requires users to put inputs on CPU before passing them
             to :func:`ChannelReducer.fit_transform`.
             Default: ``NMF``
-        **kwargs (optional): Arbitrary keyword arguments used by the specified
+        **kwargs (any, optional): Arbitrary keyword arguments used by the specified
             reduction_alg.
     """
 
@@ -72,7 +80,8 @@ def fit_transform(
         self, x: torch.Tensor, swap_2nd_and_last_dims: bool = True
     ) -> torch.Tensor:
         """
-        Perform dimensionality reduction on an input tensor.
+        Perform dimensionality reduction on an input tensor using the specified
+        ``reduction_alg``'s ``.fit_transform`` function.
 
         Args: