4320 update docstrings for gradient based saliency (#5268)

Signed-off-by: Wenqi Li <[email protected]>

Fixes #4320
it seems we also have sufficient modules to close #3542


### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [x] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

docs/source/visualize.rst

@@ -25,7 +25,15 @@ Occlusion sensitivity
 .. automodule:: monai.visualize.occlusion_sensitivity
   :members:
 
+Gradient-based saliency maps
+----------------------------
+
+.. automodule:: monai.visualize.gradient_based
+  :members:
+
+
 Utilities
 ---------
+
 .. automodule:: monai.visualize.utils
   :members:

monai/visualize/gradient_based.py

@@ -45,12 +45,29 @@ def backward(ctx, grad_output):
 
 
 class _GradReLU(torch.nn.Module):
+    """
+    A customized ReLU with the backward pass imputed for guided backpropagation (https://arxiv.org/abs/1412.6806).
+    """
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         out: torch.Tensor = _AutoGradReLU.apply(x)
         return out
 
 
 class VanillaGrad:
+    """
+    Given an input image ``x``, calling this class will perform the forward pass, then set to zero
+    all activations except one (defined by ``index``) and propagate back to the image to achieve a gradient-based
+    saliency map.
+
+    If ``index`` is None, argmax of the output logits will be used.
+
+    See also:
+
+        - Simonyan et al. Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps
+          (https://arxiv.org/abs/1312.6034)
+    """
+
     def __init__(self, model: torch.nn.Module) -> None:
         if not isinstance(model, ModelWithHooks):  # Convert to model with hooks if necessary
             self._model = ModelWithHooks(model, target_layer_names=(), register_backward=True)
@@ -83,7 +100,11 @@ def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **k
 
 class SmoothGrad(VanillaGrad):
     """
+    Compute averaged sensitivity map based on ``n_samples`` (Gaussian additive) of noisy versions
+    of the input image ``x``.
+
     See also:
+
         - Smilkov et al. SmoothGrad: removing noise by adding noise https://arxiv.org/abs/1706.03825
     """
 
@@ -126,12 +147,26 @@ def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **k
 
 
 class GuidedBackpropGrad(VanillaGrad):
+    """
+    Based on Springenberg and Dosovitskiy et al. https://arxiv.org/abs/1412.6806,
+    compute gradient-based saliency maps by backpropagating positive graidents and inputs (see ``_AutoGradReLU``).
+
+    See also:
+
+        - Springenberg and Dosovitskiy et al. Striving for Simplicity: The All Convolutional Net
+          (https://arxiv.org/abs/1412.6806)
+    """
+
     def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **kwargs) -> torch.Tensor:
         with replace_modules_temp(self.model, "relu", _GradReLU(), strict_match=False):
             return super().__call__(x, index, **kwargs)
 
 
 class GuidedBackpropSmoothGrad(SmoothGrad):
+    """
+    Compute gradient-based saliency maps based on both ``GuidedBackpropGrad`` and ``SmoothGrad``.
+    """
+
     def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **kwargs) -> torch.Tensor:
         with replace_modules_temp(self.model, "relu", _GradReLU(), strict_match=False):
             return super().__call__(x, index, **kwargs)