Rename use_autograd_hacks to sample_wise_grads_per_batch (#879)

Summary:
Pull Request resolved: #879

This diff rename use_autograd_hacks to sample_wise_grads_per_batch.

Also updated the docs for `sample_wise_grads_per_batch`.

Reviewed By: 99warriors

Differential Revision: D34506676

fbshipit-source-id: e60c2d6c09915f31406208dd0946e996b1d5719a

Author: NarineK

captum/_utils/gradient.py

@@ -772,7 +772,7 @@ def _compute_jacobian_wrt_params(
         return tuple(grads)
 
 
-def _compute_jacobian_wrt_params_autograd_hacks(
+def _compute_jacobian_wrt_params_with_sample_wise_trick(
     model: Module,
     inputs: Tuple[Any, ...],
     labels: Optional[Tensor] = None,
@@ -781,8 +781,9 @@ def _compute_jacobian_wrt_params_autograd_hacks(
 ) -> Tuple[Any, ...]:
     r"""
     Computes the Jacobian of a batch of test examples given a model, and optional
-    loss function and target labels. This method uses autograd_hacks to fully vectorize
-    the Jacobian calculation. Currently, only linear and conv2d layers are supported.
+    loss function and target labels. This method uses sample-wise gradients per
+    batch trick to fully vectorize the Jacobian calculation. Currently, only
+    linear and conv2d layers are supported.
 
     User must `add_hooks(model)` before calling this function.
 

captum/influence/_core/tracincp.py

@@ -9,7 +9,7 @@
 from captum._utils.av import AV
 from captum._utils.gradient import (
     _compute_jacobian_wrt_params,
-    _compute_jacobian_wrt_params_autograd_hacks,
+    _compute_jacobian_wrt_params_with_sample_wise_trick,
 )
 from captum.influence._core.influence import DataInfluence
 from captum.influence._utils.common import (
@@ -347,7 +347,7 @@ def __init__(
         layers: Optional[List[str]] = None,
         loss_fn: Optional[Union[Module, Callable]] = None,
         batch_size: Union[int, None] = 1,
-        use_autograd_hacks: bool = False,
+        sample_wise_grads_per_batch: bool = False,
     ) -> None:
         r"""
         Args:
@@ -396,9 +396,18 @@ def __init__(
                     `influence_src_dataset` is a Dataset. If `influence_src_dataset`
                     is a DataLoader, then `batch_size` is ignored as an argument.
                     Default: 1
-            use_autograd_hacks (bool, optional): Experimental mode that vectorize
-                    jacobian computation w.r.t parameters for a batch of inputs. Based
-                    on support in autograd_hacks.
+            sample_wise_grads_per_batch (bool, optional): PyTorch's native gradient
+                    computations w.r.t. model parameters aggregates the results for a
+                    batch and does not allow to access sample-wise gradients w.r.t.
+                    model parameters. This forces us to iterate over each sample in
+                    the batch if we want sample-wise gradients which is computationally
+                    inefficient. We offer an implementation of batch-wise gradient
+                    computations w.r.t. to model parameters which is computationally
+                    more efficient. This implementation can be enabled by setting the
+                    `sample_wise_grad_per_batch` argument to `True`. Note that our
+                    current implementation enables batch-wise gradient computations
+                    only for a limited number of PyTorch nn.Modules: Conv2D and Linear.
+                    This list will be expanded in the near future.
                     Default: False
         """
 
@@ -412,10 +421,10 @@ def __init__(
             batch_size,
         )
 
-        self.use_autograd_hacks = use_autograd_hacks
+        self.sample_wise_grads_per_batch = sample_wise_grads_per_batch
 
         if (
-            self.use_autograd_hacks
+            self.sample_wise_grads_per_batch
             and isinstance(loss_fn, Module)  # TODO: allow loss_fn to be Callable
             and hasattr(loss_fn, "reduction")
         ):
@@ -644,8 +653,8 @@ def _basic_computation_tracincp(
             targets (tensor or None): If computing influence scores on a loss function,
                     these are the labels corresponding to the batch `inputs`.
         """
-        if self.use_autograd_hacks:
-            return _compute_jacobian_wrt_params_autograd_hacks(
+        if self.sample_wise_grads_per_batch:
+            return _compute_jacobian_wrt_params_with_sample_wise_trick(
                 self.model,
                 inputs,
                 targets,

tests/influence/_core/test_tracincp.py

@@ -7,10 +7,10 @@
     _TestTracInRegression20DCheckIdx,
     _TestTracInXORCheckIdx,
     _TestTracInIdentityRegressionCheckIdx,
-    _TestTracInRegression1DCheckAutogradHacks,
-    _TestTracInRegression20DCheckAutogradHacks,
-    _TestTracInXORCheckAutogradHacks,
-    _TestTracInIdentityRegressionCheckAutogradHacks,
+    _TestTracInRegression1DCheckSampleWiseTrick,
+    _TestTracInRegression20DCheckSampleWiseTrick,
+    _TestTracInXORCheckSampleWiseTrick,
+    _TestTracInIdentityRegressionCheckSampleWiseTrick,
     _TestTracInRegression1DNumerical,
     _TestTracInGetKMostInfluential,
     _TestTracInSelfInfluence,
@@ -38,29 +38,29 @@ def setUp(self):
                 tmpdir,
                 batch_size=batch_size,
                 loss_fn=loss_fn,
-                use_autograd_hacks=False,
+                sample_wise_grads_per_batch=False,
             )
         )
         super(TestTracInCP, self).setUp()
 
 
-class TestTracInCPCheckAutogradHacks(
-    _TestTracInRegression1DCheckAutogradHacks,
-    _TestTracInRegression20DCheckAutogradHacks,
-    _TestTracInXORCheckAutogradHacks,
-    _TestTracInIdentityRegressionCheckAutogradHacks,
+class TestTracInCPCheckSampleWiseTrick(
+    _TestTracInRegression1DCheckSampleWiseTrick,
+    _TestTracInRegression20DCheckSampleWiseTrick,
+    _TestTracInXORCheckSampleWiseTrick,
+    _TestTracInIdentityRegressionCheckSampleWiseTrick,
     BaseTest,
 ):
     def setUp(self):
         self.tracin_constructor = (
-            lambda net, dataset, tmpdir, batch_size, loss_fn, use_autograd_hacks: (
+            lambda net, dataset, tmpdir, batch_size, loss_fn, sample_wise_trick: (
                 TracInCP(
                     net,
                     dataset,
                     tmpdir,
                     batch_size=batch_size,
                     loss_fn=loss_fn,
-                    use_autograd_hacks=use_autograd_hacks,
+                    sample_wise_grads_per_batch=sample_wise_trick,
                 )
             )
         )
@@ -82,6 +82,6 @@ def setUp(self):
                 tmpdir,
                 batch_size=batch_size,
                 loss_fn=loss_fn,
-                use_autograd_hacks=True,
+                sample_wise_grads_per_batch=True,
             )
         )

tests/influence/_utils/common.py

@@ -215,7 +215,7 @@ def _test_tracin_regression(self, features: int, mode: int) -> None:
                 for i in range(len(idx)):
                     self.assertTrue(isSorted(idx[i]))
 
-            if mode == "check_autograd_hacks":
+            if mode == "sample_wise_trick":
 
                 criterion = nn.MSELoss(reduction="none")
 
@@ -228,39 +228,47 @@ def _test_tracin_regression(self, features: int, mode: int) -> None:
                     False,
                 )
 
-                # With autograd hacks
+                # With sample-wise trick
                 criterion = nn.MSELoss(reduction="sum")
-                tracin_hack = self.tracin_constructor(
+                tracin_sample_wise_trick = self.tracin_constructor(
                     net, dataset, tmpdir, batch_size, criterion, True
                 )
 
                 train_scores = tracin.influence(train_inputs, train_labels)
-                train_scores_hack = tracin_hack.influence(train_inputs, train_labels)
-                assertTensorAlmostEqual(self, train_scores, train_scores_hack)
+                train_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
+                    train_inputs, train_labels
+                )
+                assertTensorAlmostEqual(
+                    self, train_scores, train_scores_sample_wise_trick
+                )
 
                 test_scores = tracin.influence(test_inputs, test_labels)
-                test_scores_hack = tracin_hack.influence(test_inputs, test_labels)
-                assertTensorAlmostEqual(self, test_scores, test_scores_hack)
+                test_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
+                    test_inputs, test_labels
+                )
+                assertTensorAlmostEqual(
+                    self, test_scores, test_scores_sample_wise_trick
+                )
 
 
 class _TestTracInRegression1DCheckIdx(_TestTracInRegression):
     def test_tracin_regression_1D_check_idx(self):
         self._test_tracin_regression(1, "check_idx")
 
 
-class _TestTracInRegression1DCheckAutogradHacks(_TestTracInRegression):
-    def test_tracin_regression_1D_check_autograd_hacks(self):
-        self._test_tracin_regression(1, "check_autograd_hacks")
+class _TestTracInRegression1DCheckSampleWiseTrick(_TestTracInRegression):
+    def test_tracin_regression_1D_check_sample_wise_trick(self):
+        self._test_tracin_regression(1, "sample_wise_trick")
 
 
 class _TestTracInRegression20DCheckIdx(_TestTracInRegression):
     def test_tracin_regression_20D_check_idx(self):
         self._test_tracin_regression(20, "check_idx")
 
 
-class _TestTracInRegression20DCheckAutogradHacks(_TestTracInRegression):
-    def test_tracin_regression_20D_check_autograd_hacks(self):
-        self._test_tracin_regression(20, "check_autograd_hacks")
+class _TestTracInRegression20DCheckSampleWiseTrick(_TestTracInRegression):
+    def test_tracin_regression_20D_check_sample_wise_trick(self):
+        self._test_tracin_regression(20, "sample_wise_tricksample_wise_trick")
 
 
 class _TestTracInXOR:
@@ -434,7 +442,7 @@ def _test_tracin_xor(self, mode) -> None:
                     influence_labels = dataset.labels[idx[i][0:5], 0]
                     self.assertTrue(torch.all(testlabels[i, 0] == influence_labels))
 
-            if mode == "check_autograd_hacks":
+            if mode == "sample_wise_trick":
 
                 criterion = nn.MSELoss(reduction="none")
 
@@ -447,9 +455,9 @@ def _test_tracin_xor(self, mode) -> None:
                     False,
                 )
 
-                # With autograd hacks
+                # With sample-wise trick
                 criterion = nn.MSELoss(reduction="sum")
-                tracin_hack = self.tracin_constructor(
+                tracin_sample_wise_trick = self.tracin_constructor(
                     net,
                     dataset,
                     tmpdir,
@@ -459,18 +467,22 @@ def _test_tracin_xor(self, mode) -> None:
                 )
 
                 test_scores = tracin.influence(testset, testlabels)
-                test_scores_hack = tracin_hack.influence(testset, testlabels)
-                assertTensorAlmostEqual(self, test_scores, test_scores_hack)
+                test_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
+                    testset, testlabels
+                )
+                assertTensorAlmostEqual(
+                    self, test_scores, test_scores_sample_wise_trick
+                )
 
 
 class _TestTracInXORCheckIdx(_TestTracInXOR):
     def test_tracin_xor_check_idx(self):
         self._test_tracin_xor("check_idx")
 
 
-class _TestTracInXORCheckAutogradHacks(_TestTracInXOR):
-    def test_tracin_xor_check_autograd_hacks(self):
-        self._test_tracin_xor("check_autograd_hacks")
+class _TestTracInXORCheckSampleWiseTrick(_TestTracInXOR):
+    def test_tracin_xor_check_sample_wise_trick(self):
+        self._test_tracin_xor("sample_wise_trick")
 
 
 class _TestTracInIdentityRegression:
@@ -537,7 +549,7 @@ def _test_tracin_identity_regression(self, mode) -> None:
                 for i in range(len(idx)):
                     self.assertEqual(idx[i][0], i)
 
-            if mode == "check_autograd_hacks":
+            if mode == "sample_wise_trick":
 
                 criterion = nn.MSELoss(reduction="none")
 
@@ -550,9 +562,9 @@ def _test_tracin_identity_regression(self, mode) -> None:
                     False,
                 )
 
-                # With autograd hacks
+                # With sample-wise trick
                 criterion = nn.MSELoss(reduction="sum")
-                tracin_hack = self.tracin_constructor(
+                tracin_sample_wise_trick = self.tracin_constructor(
                     net,
                     dataset,
                     tmpdir,
@@ -562,18 +574,22 @@ def _test_tracin_identity_regression(self, mode) -> None:
                 )
 
                 train_scores = tracin.influence(train_inputs, train_labels)
-                train_scores_hack = tracin_hack.influence(train_inputs, train_labels)
-                assertTensorAlmostEqual(self, train_scores, train_scores_hack)
+                train_scores_tracin_sample_wise_trick = (
+                    tracin_sample_wise_trick.influence(train_inputs, train_labels)
+                )
+                assertTensorAlmostEqual(
+                    self, train_scores, train_scores_tracin_sample_wise_trick
+                )
 
 
 class _TestTracInIdentityRegressionCheckIdx(_TestTracInIdentityRegression):
     def test_tracin_identity_regression_check_idx(self):
         self._test_tracin_identity_regression("check_idx")
 
 
-class _TestTracInIdentityRegressionCheckAutogradHacks(_TestTracInIdentityRegression):
-    def test_tracin_identity_regression_check_autograd_hacks(self):
-        self._test_tracin_identity_regression("check_autograd_hacks")
+class _TestTracInIdentityRegressionCheckSampleWiseTrick(_TestTracInIdentityRegression):
+    def test_tracin_identity_regression_check_sample_wise_trick(self):
+        self._test_tracin_identity_regression("sample_wise_trick")
 
 
 class _TestTracInRandomProjectionRegression: