Add capability to pass grad_kwargs for grad_cam, internal_influence, layer_conductance, layer_deep_lift, layer_gradient_shap, and neuron_conductance (#1294)

Summary:
Pull Request resolved: #1294

Extension of D57756842, where the torch.autograd.grad arguments can be passed into the following classes

LayerGradCam
InternalInfluence
LayerConductance
LayerDeepLift
LayerGradientShap
NeuronConductance

Reviewed By: cyrjano

Differential Revision: D58208128

fbshipit-source-id: ec8449e5c51d000c70c4b691858b96a0963728ff

Author: jjuncho

captum/attr/_core/layer/grad_cam.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -83,6 +83,7 @@ def attribute(
         attribute_to_layer_input: bool = False,
         relu_attributions: bool = False,
         attr_dim_summation: bool = True,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]:
         r"""
         Args:
@@ -154,6 +155,9 @@ def attribute(
                         sum attributions along dimension 1 (usually channel).
                         The default (True) means to sum along dimension 1.
                         Default: True
+            grad_kwargs (Dict[str, Any], optional): Additional keyword
+                        arguments for torch.autograd.grad.
+                        Default: None
 
         Returns:
             *Tensor* or *tuple[Tensor, ...]* of **attributions**:
@@ -200,6 +204,7 @@ def attribute(
             additional_forward_args,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_layer_input,
+            grad_kwargs=grad_kwargs,
         )
 
         summed_grads = tuple(

captum/attr/_core/layer/internal_influence.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -74,6 +74,7 @@ def attribute(
         method: str = "gausslegendre",
         internal_batch_size: Union[None, int] = None,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]:
         r"""
         Args:
@@ -185,6 +186,9 @@ def attribute(
                         attribute to the input or output, is a single tensor.
                         Support for multiple tensors will be added later.
                         Default: False
+            grad_kwargs (Dict[str, Any], optional): Additional keyword
+                        arguments for torch.autograd.grad.
+                        Default: None
 
         Returns:
             *Tensor* or *tuple[Tensor, ...]* of **attributions**:
@@ -236,6 +240,7 @@ def attribute(
                 n_steps=n_steps,
                 method=method,
                 attribute_to_layer_input=attribute_to_layer_input,
+                grad_kwargs=grad_kwargs,
             )
 
         return attrs
@@ -250,6 +255,7 @@ def _attribute(
         method: str = "gausslegendre",
         attribute_to_layer_input: bool = False,
         step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]:
         if step_sizes_and_alphas is None:
             # retrieve step size and scaling factor for specified approximation method
@@ -290,6 +296,7 @@ def _attribute(
             additional_forward_args=input_additional_args,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_layer_input,
+            grad_kwargs=grad_kwargs,
         )
         # flattening grads so that we can multiply it with step-size
         # calling contiguous to avoid `memory whole` problems

captum/attr/_core/layer/layer_conductance.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import typing
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -82,6 +82,7 @@ def attribute(
         *,
         return_convergence_delta: Literal[True],
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]: ...
 
     @typing.overload
@@ -96,6 +97,7 @@ def attribute(
         internal_batch_size: Union[None, int] = None,
         return_convergence_delta: Literal[False] = False,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]: ...
 
     @log_usage()
@@ -112,6 +114,7 @@ def attribute(
         internal_batch_size: Union[None, int] = None,
         return_convergence_delta: bool = False,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[
         Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
     ]:
@@ -230,6 +233,9 @@ def attribute(
                         attribute to the input or output, is a single tensor.
                         Support for multiple tensors will be added later.
                         Default: False
+            grad_kwargs (Dict[str, Any], optional): Additional keyword
+                        arguments for torch.autograd.grad.
+                        Default: None
 
         Returns:
             **attributions** or 2-element tuple of **attributions**, **delta**:
@@ -322,6 +328,7 @@ def _attribute(
         method: str = "gausslegendre",
         attribute_to_layer_input: bool = False,
         step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]:
         num_examples = inputs[0].shape[0]
         if step_sizes_and_alphas is None:
@@ -366,6 +373,7 @@ def _attribute(
             target_ind=expanded_target,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_layer_input,
+            grad_kwargs=grad_kwargs,
         )
 
         # Compute differences between consecutive evaluations of layer_eval.

captum/attr/_core/layer/layer_deep_lift.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import typing
-from typing import Any, Callable, cast, Sequence, Tuple, Union
+from typing import Any, Callable, cast, Dict, Optional, Sequence, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -108,6 +108,7 @@ def attribute(
         return_convergence_delta: Literal[False] = False,
         attribute_to_layer_input: bool = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]: ...
 
     @typing.overload
@@ -121,6 +122,7 @@ def attribute(
         return_convergence_delta: Literal[True],
         attribute_to_layer_input: bool = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]: ...
 
     @log_usage()
@@ -133,6 +135,7 @@ def attribute(
         return_convergence_delta: bool = False,
         attribute_to_layer_input: bool = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[
         Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
     ]:
@@ -248,6 +251,9 @@ def attribute(
                         `custom_attribution_func` returns a tuple of attribution
                         tensors that have the same length as the `inputs`.
                         Default: None
+            grad_kwargs (Dict[str, Any], optional): Additional keyword
+                        arguments for torch.autograd.grad.
+                        Default: None
 
         Returns:
             **attributions** or 2-element tuple of **attributions**, **delta**:
@@ -274,6 +280,7 @@ def attribute(
                 it is not guaranteed and depends on the specifics of the
                 `custom_attribution_func`.
 
+
         Examples::
 
             >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
@@ -326,6 +333,7 @@ def chunk_output_fn(out: TensorOrTupleOfTensorsGeneric) -> Sequence:
                 inputs,
                 attribute_to_layer_input=attribute_to_layer_input,
                 output_fn=lambda out: chunk_output_fn(out),
+                grad_kwargs=grad_kwargs,
             )
 
             attr_inputs = tuple(map(lambda attr: attr[0], attrs))

captum/attr/_core/layer/layer_gradient_shap.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import typing
-from typing import Any, Callable, cast, List, Tuple, Union
+from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -242,6 +242,7 @@ def attribute(
                         attribute to the input or output, is a single tensor.
                         Support for multiple tensors will be added later.
                         Default: False
+
         Returns:
             **attributions** or 2-element tuple of **attributions**, **delta**:
             - **attributions** (*Tensor* or *tuple[Tensor, ...]*):
@@ -375,6 +376,7 @@ def attribute(
         additional_forward_args: Any = None,
         return_convergence_delta: Literal[False] = False,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...]]: ...
 
     @typing.overload
@@ -387,6 +389,7 @@ def attribute(
         *,
         return_convergence_delta: Literal[True],
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]: ...
 
     @log_usage()
@@ -398,6 +401,7 @@ def attribute(  # type: ignore
         additional_forward_args: Any = None,
         return_convergence_delta: bool = False,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[
         Tensor, Tuple[Tensor, ...], Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]
     ]:
@@ -420,6 +424,7 @@ def attribute(  # type: ignore
             additional_forward_args,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_layer_input,
+            grad_kwargs=grad_kwargs,
         )
 
         attr_baselines = _forward_layer_eval(

captum/attr/_core/layer/layer_gradient_x_activation.py

@@ -133,8 +133,9 @@ def attribute(
                         layer input, otherwise it will be computed with respect
                         to layer output.
                         Default: False
-            grad_kwargs: Additional keyword arguments for torch.autograd.grad
-
+            grad_kwargs (Dict[str, Any], optional): Additional keyword
+                        arguments for torch.autograd.grad.
+                        Default: None
         Returns:
             *Tensor* or *tuple[Tensor, ...]* or list of **attributions**:
             - **attributions** (*Tensor*, *tuple[Tensor, ...]*, or *list*):

captum/attr/_core/neuron/neuron_conductance.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import warnings
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -99,6 +99,7 @@ def attribute(
         method: str = "riemann_trapezoid",
         internal_batch_size: Union[None, int] = None,
         attribute_to_neuron_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> TensorOrTupleOfTensorsGeneric:
         r"""
         Args:
@@ -311,6 +312,7 @@ def attribute(
                 n_steps=n_steps,
                 method=method,
                 attribute_to_neuron_input=attribute_to_neuron_input,
+                grad_kwargs=grad_kwargs,
             )
         return _format_output(is_inputs_tuple, attrs)
 
@@ -325,6 +327,7 @@ def _attribute(
         method: str = "riemann_trapezoid",
         attribute_to_neuron_input: bool = False,
         step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Tensor, ...]:
 
         num_examples = inputs[0].shape[0]
@@ -371,6 +374,7 @@ def _attribute(
             gradient_neuron_selector=neuron_selector,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_neuron_input,
+            grad_kwargs=grad_kwargs,
         )
 
         mid_grads = _verify_select_neuron(layer_gradients, neuron_selector)

tests/attr/layer/test_grad_cam.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 
 import unittest
-from typing import Any, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 from captum._utils.typing import TensorLikeList
 from captum.attr._core.layer.grad_cam import LayerGradCam
+from packaging import version
 from tests.helpers import BaseTest
 from tests.helpers.basic import assertTensorTuplesAlmostEqual
 from tests.helpers.basic_models import (
@@ -119,6 +120,7 @@ def _grad_cam_test_assert(
         attribute_to_layer_input: bool = False,
         relu_attributions: bool = False,
         attr_dim_summation: bool = True,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
         layer_gc = LayerGradCam(model, target_layer)
         self.assertFalse(layer_gc.multiplies_by_inputs)
@@ -129,11 +131,30 @@ def _grad_cam_test_assert(
             attribute_to_layer_input=attribute_to_layer_input,
             relu_attributions=relu_attributions,
             attr_dim_summation=attr_dim_summation,
+            grad_kwargs=grad_kwargs,
         )
         assertTensorTuplesAlmostEqual(
             self, attributions, expected_activation, delta=0.01
         )
 
+    def test_relu_gradcam_with_unused_layer(self) -> None:
+        if version.parse(torch.__version__) < version.parse("2.1.0"):
+            raise unittest.SkipTest(
+                "Skipping unused layed gradient test since it is not supported "
+                "by torch version < 2.1"
+            )
+        net = BasicModel_MultiLayer(multi_input_module=True)
+        inp = torch.tensor([[0.0, 6.0, 0.0]], requires_grad=True)
+        gradcam = LayerGradCam(net, net.relu)
+        attributions = gradcam.attribute(
+            inputs=inp,
+            target=0,
+            grad_kwargs={"materialize_grads": True},
+        )
+        self.assertEqual(len(attributions), 1)
+        self.assertEqual(list(attributions[0].shape), [1])
+        self.assertAlmostEqual(attributions[0].sum(), 0)
+
 
 if __name__ == "__main__":
     unittest.main()