Add capability to pass additional grad_kwargs for LayerGradientXActivation (#1286)

Summary: Pull Request resolved: #1286

Differential Revision: D57756842

Author: yucu

captum/_utils/gradient.py

@@ -485,6 +485,7 @@ def compute_layer_gradients_and_eval(
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     output_fn: Union[None, Callable] = None,
+    grad_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...]]: ...
 
 
@@ -499,6 +500,7 @@ def compute_layer_gradients_and_eval(
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     output_fn: Union[None, Callable] = None,
+    grad_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Tuple[List[Tuple[Tensor, ...]], List[Tuple[Tensor, ...]]]: ...
 
 
@@ -513,6 +515,7 @@ def compute_layer_gradients_and_eval(
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     output_fn: Union[None, Callable] = None,
+    grad_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]]: ...
 
 
@@ -528,6 +531,7 @@ def compute_layer_gradients_and_eval(
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     output_fn: Union[None, Callable] = None,
+    grad_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Union[
     Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]],
     Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...]],
@@ -572,6 +576,7 @@ def compute_layer_gradients_and_eval(
         args:       Additional input arguments that forward function requires.
                     It takes an empty tuple (no additional arguments) if no
                     additional arguments are required
+        grad_kwargs: Additional keyword arguments for torch.autograd.grad
 
 
     Returns:
@@ -640,7 +645,11 @@ def compute_layer_gradients_and_eval(
             for device_id in key_list
             for layer_tensor in saved_layer[single_layer][device_id]
         )
-        saved_grads = torch.autograd.grad(torch.unbind(output), grad_inputs)
+        saved_grads = torch.autograd.grad(
+            outputs=torch.unbind(output),
+            inputs=grad_inputs,
+            **grad_kwargs or {},
+        )
 
         offset = 0
         all_grads: List[Tuple[Tensor, ...]] = []

captum/attr/_core/layer/layer_gradient_x_activation.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
 
 from captum._utils.common import (
     _format_additional_forward_args,
@@ -76,6 +76,7 @@ def attribute(
         target: TargetType = None,
         additional_forward_args: Any = None,
         attribute_to_layer_input: bool = False,
+        grad_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[Tensor, Tuple[Tensor, ...], List[Union[Tensor, Tuple[Tensor, ...]]]]:
         r"""
         Args:
@@ -132,6 +133,7 @@ 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
 
         Returns:
             *Tensor* or *tuple[Tensor, ...]* or list of **attributions**:
@@ -175,6 +177,7 @@ def attribute(
             additional_forward_args,
             device_ids=self.device_ids,
             attribute_to_layer_input=attribute_to_layer_input,
+            grad_kwargs=grad_kwargs,
         )
         if isinstance(self.layer, Module):
             return _format_output(

tests/attr/layer/test_layer_gradient_x_activation.py

@@ -129,6 +129,18 @@ def test_gradient_activation_embedding_no_grad(self) -> None:
                 list(layer_act.attribute(inputs=(input1, input2)).shape), [4, 100]
             )
 
+    def test_simple_multi_gradient_activation_with_unused_layer(self) -> None:
+        model = BasicModel_MultiLayer(multi_input_module=True)
+        test_input1 = torch.tensor([[3.0, 4.0, 0.0]], requires_grad=True)
+        # test_input2 = torch.tensor([[0.0, 4.0, 5.0]], requires_grad=True)
+        layer_act = LayerGradientXActivation(model, [model.linear1, model.relu])
+        attributions = layer_act.attribute(
+            inputs=test_input1, target=0, grad_kwargs={"materialize_grads": True}
+        )
+        self.assertEqual(len(attributions), 2)
+        self.assertEqual(list(attributions[0].shape), [1, 4])
+        self.assertEqual(list(attributions[1].shape), [1, 4])
+
     def _layer_activation_test_assert(
         self,
         model: Module,

tests/helpers/basic_models.py

@@ -407,6 +407,10 @@ def forward(
         if self.multi_input_module:
             relu_out1, relu_out2 = self.multi_relu(lin1_out, self.linear1_alt(input))
             relu_out = relu_out1 + relu_out2
+            # relu is not used when multi_input_module set to True,
+            # so this is to set an unsued layer intentionally for testing
+            # and it won't be part of return
+            self.relu(lin1_out)
         else:
             relu_out = self.relu(lin1_out)
         lin2_out = self.linear2(relu_out)

tests/utils/test_gradient.py

@@ -243,3 +243,20 @@ def test_layer_gradient_output(self) -> None:
         )
         assertTensorAlmostEqual(self, grads[0], [[0.0, 1.0]], delta=0.01, mode="max")
         assertTensorAlmostEqual(self, eval[0], [[26.0, 28.0]], delta=0.01, mode="max")
+
+    def test_layer_gradient_unused_layer(self) -> None:
+        model = BasicModel_MultiLayer(multi_input_module=True)
+        input = torch.tensor([[5.0, 2.0, 1.0]], requires_grad=True)
+        grads, eval = compute_layer_gradients_and_eval(
+            model,
+            [model.linear1, model.relu],
+            input,
+            target_ind=1,
+            grad_kwargs={"materialize_grads": True},
+        )
+        assertTensorAlmostEqual(
+            self, grads[0][0], [[0.0, 1.0, 1.0, 1.0]], delta=0, mode="max"
+        )
+        assertTensorAlmostEqual(
+            self, eval[0][0], [[-2.0, 9.0, 9.0, 9.0]], delta=0, mode="max"
+        )