update tracin influence API (#1072)

Summary:
Pull Request resolved: #1072

This diff changes the API for implementations of `TracInCPBase` as discussed in https://fb.quip.com/JbpnAiWluZmI.  In particular, the arguments representing test data of the `influence` method are changed from `inputs: Tuple, targets: Optional[Tensor]` to `inputs: Union[Tuple[Any], DataLoader]`, which is either a single batch, or a dataloader yielding batches.  In both cases, `model(*batch)` is assumed to produce the predictions for a batch, and `batch[-1]` is assumed to be the labels for a batch. This is the same format assumed of the batches yielded by `train_dataloader`.

We make this change for 2 reasons
- it unifies the assumptions made of the test data and the assumptions made of the training data
- for some implementations, we want to allow the test data to be represented by a dataloader.  with the old API, there was no clean way to allow both a single as well as a dataloader to be passed in, since a batch required 2 arguments, but a dataloader only requires 1.

For now, all implementations only allow `inputs` to be a tuple (and not a dataloader).  This is okay due to inheritance rules.  Later on, we will allow some implementations (i.e. `TracInCP`) to accept a dataloader as `inputs`.

Other changes:
- changes to make documentation.  for example, documentation in `TracInCPBase.influence` now refers to the "test dataset" instead of test batch.
- the `unpack_inputs` argument is no longer needed for the `influence` methods, and is removed
- the usage of `influence` in all the tests is changed to match new API.
- signature of helper methods `_influence_batch_tracincp` and `_influence_batch_tracincp_fast` are changed to match new representation of batches.

Reviewed By: cyrjano

Differential Revision: D41324297

fbshipit-source-id: 9fe108de1a6789d461c19d71b724cd18bbcffbd9

Author: Fulton Wang

captum/influence/_core/influence.py

@@ -12,21 +12,19 @@ class DataInfluence(ABC):
     An abstract class to define model data influence skeleton.
     """
 
-    def __init_(
-        self, model: Module, influence_src_dataset: Dataset, **kwargs: Any
-    ) -> None:
+    def __init_(self, model: Module, train_dataset: Dataset, **kwargs: Any) -> None:
         r"""
         Args:
             model (torch.nn.Module): An instance of pytorch model.
-            influence_src_dataset (torch.utils.data.Dataset): PyTorch Dataset that is
+            train_dataset (torch.utils.data.Dataset): PyTorch Dataset that is
                     used to create a PyTorch Dataloader to iterate over the dataset and
                     its labels. This is the dataset for which we will be seeking for
                     influential instances. In most cases this is the training dataset.
             **kwargs: Additional key-value arguments that are necessary for specific
                     implementation of `DataInfluence` abstract class.
         """
         self.model = model
-        self.influence_src_dataset = influence_src_dataset
+        self.train_dataset = train_dataset
 
     @abstractmethod
     def influence(self, inputs: Any = None, **kwargs: Any) -> Any:

captum/influence/_core/tracincp.py

@@ -191,7 +191,6 @@ def _get_k_most_influential_helper(
     influence_src_dataloader: DataLoader,
     influence_batch_fn: Callable,
     inputs: Tuple[Any, ...],
-    targets: Optional[Tensor],
     k: int = 5,
     proponents: bool = True,
     show_progress: bool = False,
@@ -206,13 +205,12 @@ def _get_k_most_influential_helper(
         influence_src_dataloader (DataLoader): The DataLoader, representing training
                 data, for which we want to compute proponents / opponents.
         influence_batch_fn (Callable): A callable that will be called via
-                `influence_batch_fn(inputs, targets, batch)`, where `batch` is a batch
+                `influence_batch_fn(inputs, batch)`, where `batch` is a batch
                 in the `influence_src_dataloader` argument.
-        inputs (tuple[Any, ...]): A batch of examples. Does not represent labels,
-                which are passed as `targets`.
-        targets (Tensor, optional): If computing TracIn scores on a loss function,
-                these are the labels corresponding to the batch `inputs`.
-                Default: None
+        inputs (tuple[Any, ...]): This argument represents the test batch, and is a
+                single tuple of any, where the last element is assumed to be the labels
+                for the batch. That is, `model(*batch[0:-1])` produces the output for
+                `model`, and `batch[-1]` are the labels, if any.
         k (int, optional): The number of proponents or opponents to return per test
                 instance.
                 Default: 5
@@ -274,7 +272,7 @@ def _get_k_most_influential_helper(
     for batch in influence_src_dataloader:
 
         # calculate tracin_scores for the batch
-        batch_tracin_scores = influence_batch_fn(inputs, targets, batch)
+        batch_tracin_scores = influence_batch_fn(inputs, batch)
         batch_tracin_scores *= multiplier
 
         # get the top-k indices and tracin_scores for the batch

captum/influence/_core/tracincp_fast_rand_proj.py

@@ -10,6 +10,7 @@
 from parameterized import parameterized
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.influence._utils.common import (
+    _format_batch_into_tuple,
     build_test_name_func,
     DataInfluenceConstructor,
     get_random_model_and_data,
@@ -76,7 +77,8 @@ def test_tracin_dataloader(
             )
 
             train_scores = tracin.influence(
-                test_samples, test_labels, k=None, unpack_inputs=unpack_inputs
+                _format_batch_into_tuple(test_samples, test_labels, unpack_inputs),
+                k=None,
             )
 
             tracin_dataloader = tracin_constructor(
@@ -88,7 +90,8 @@ def test_tracin_dataloader(
             )
 
             train_scores_dataloader = tracin_dataloader.influence(
-                test_samples, test_labels, k=None, unpack_inputs=unpack_inputs
+                _format_batch_into_tuple(test_samples, test_labels, unpack_inputs),
+                k=None,
             )
 
             assertTensorAlmostEqual(

captum/influence/_utils/common.py

@@ -12,6 +12,7 @@
 from parameterized import parameterized
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.influence._utils.common import (
+    _format_batch_into_tuple,
     build_test_name_func,
     DataInfluenceConstructor,
     get_random_model_and_data,
@@ -224,25 +225,13 @@ def test_tracin_intermediate_quantities_consistent(
             )
 
             # compute influence scores without using `compute_intermediate_quantities`
+            test_batch = _format_batch_into_tuple(
+                test_features, test_labels, unpack_inputs
+            )
             scores = tracin.influence(
-                test_features, test_labels, unpack_inputs=unpack_inputs
+                test_batch,
             )
 
-            # compute influence scores using `compute_intermediate_quantities`
-            # we combine `test_features` and `test_labels` into a single tuple
-            # `test_batch` to pass to the model, with the assumption that
-            # `model(test_batch[0:-1]` produces the predictions, and `test_batch[-1]`
-            # are the labels.  We do this due to the assumptions made by the
-            # `compute_intermediate_quantities` method. Therefore, how we
-            # form `test_batch` depends on whether `unpack_inputs` is True or False
-            if not unpack_inputs:
-                # `test_features` is a Tensor
-                test_batch = (test_features, test_labels)
-            else:
-                # `test_features` is a tuple, so we unpack it to place in tuple,
-                # along with `test_labels`
-                test_batch = (*test_features, test_labels)  # type: ignore[assignment]
-
             # the influence score is the dot product of intermediate quantities
             intermediate_quantities_scores = torch.matmul(
                 intermediate_quantities_tracin.compute_intermediate_quantities(

tests/influence/_core/test_dataloader.py

@@ -8,6 +8,7 @@
 from parameterized import parameterized
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.influence._utils.common import (
+    _format_batch_into_tuple,
     build_test_name_func,
     DataInfluenceConstructor,
     get_random_model_and_data,
@@ -107,15 +108,14 @@ def test_tracin_k_most_influential(
             )
 
             train_scores = tracin.influence(
-                test_samples, test_labels, k=None, unpack_inputs=unpack_inputs
+                _format_batch_into_tuple(test_samples, test_labels, unpack_inputs),
+                k=None,
             )
             sort_idx = torch.argsort(train_scores, dim=1, descending=proponents)[:, 0:k]
             idx, _train_scores = tracin.influence(
-                test_samples,
-                test_labels,
+                _format_batch_into_tuple(test_samples, test_labels, unpack_inputs),
                 k=k,
                 proponents=proponents,
-                unpack_inputs=unpack_inputs,
             )
             for i in range(len(idx)):
                 # check that idx[i] is correct

tests/influence/_core/test_tracin_intermediate_quantities.py

@@ -183,19 +183,19 @@ def test_tracin_regression(
                     criterion,
                 )
 
-                train_scores = tracin.influence(train_inputs, train_labels)
+                train_scores = tracin.influence((train_inputs, train_labels))
                 idx, _ = tracin.influence(
-                    train_inputs, train_labels, k=len(dataset), proponents=True
+                    (train_inputs, train_labels), k=len(dataset), proponents=True
                 )
                 # check that top influence is one with maximal value
                 # (and hence gradient)
                 for i in range(len(idx)):
                     self.assertEqual(idx[i][0], 15)
 
                 # check influence scores of test data
-                test_scores = tracin.influence(test_inputs, test_labels)
+                test_scores = tracin.influence((test_inputs, test_labels))
                 idx, _ = tracin.influence(
-                    test_inputs, test_labels, k=len(test_inputs), proponents=True
+                    (test_inputs, test_labels), k=len(test_inputs), proponents=True
                 )
                 # check that top influence is one with maximal value
                 # (and hence gradient)
@@ -226,17 +226,17 @@ def test_tracin_regression(
                     sample_wise_grads_per_batch=True,
                 )
 
-                train_scores = tracin.influence(train_inputs, train_labels)
+                train_scores = tracin.influence((train_inputs, train_labels))
                 train_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
-                    train_inputs, train_labels
+                    (train_inputs, train_labels)
                 )
                 assertTensorAlmostEqual(
                     self, train_scores, train_scores_sample_wise_trick
                 )
 
-                test_scores = tracin.influence(test_inputs, test_labels)
+                test_scores = tracin.influence((test_inputs, test_labels))
                 test_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
-                    test_inputs, test_labels
+                    (test_inputs, test_labels)
                 )
                 assertTensorAlmostEqual(
                     self, test_scores, test_scores_sample_wise_trick
@@ -288,7 +288,7 @@ def test_tracin_regression_1D_numerical(
                 criterion,
             )
 
-            train_scores = tracin.influence(train_inputs, train_labels, k=None)
+            train_scores = tracin.influence((train_inputs, train_labels), k=None)
 
             r"""
             Derivation for gradient / resulting TracIn score:
@@ -382,9 +382,9 @@ def test_tracin_identity_regression(
 
                 # check influence scores of training data
 
-                train_scores = tracin.influence(train_inputs, train_labels)
+                train_scores = tracin.influence((train_inputs, train_labels))
                 idx, _ = tracin.influence(
-                    train_inputs, train_labels, k=len(dataset), proponents=True
+                    (train_inputs, train_labels), k=len(dataset), proponents=True
                 )
 
                 # check that top influence for an instance is itself
@@ -415,9 +415,9 @@ def test_tracin_identity_regression(
                     sample_wise_grads_per_batch=True,
                 )
 
-                train_scores = tracin.influence(train_inputs, train_labels)
+                train_scores = tracin.influence((train_inputs, train_labels))
                 train_scores_tracin_sample_wise_trick = (
-                    tracin_sample_wise_trick.influence(train_inputs, train_labels)
+                    tracin_sample_wise_trick.influence((train_inputs, train_labels))
                 )
                 assertTensorAlmostEqual(
                     self, train_scores, train_scores_tracin_sample_wise_trick
@@ -496,5 +496,5 @@ def test_loss_fn(input, target):
             )
 
             # check influence scores of training data. they should all be 0
-            train_scores = tracin.influence(train_inputs, train_labels, k=None)
+            train_scores = tracin.influence((train_inputs, train_labels), k=None)
             assertTensorAlmostEqual(self, train_scores, torch.zeros(train_scores.shape))

tests/influence/_core/test_tracin_k_most_influential.py

@@ -8,6 +8,7 @@
 from parameterized import parameterized
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.influence._utils.common import (
+    _format_batch_into_tuple,
     build_test_name_func,
     DataInfluenceConstructor,
     get_random_model_and_data,
@@ -108,10 +109,10 @@ def test_tracin_self_influence(
                 criterion,
             )
             train_scores = tracin.influence(
-                train_dataset.samples,
-                train_dataset.labels,
+                _format_batch_into_tuple(
+                    train_dataset.samples, train_dataset.labels, unpack_inputs
+                ),
                 k=None,
-                unpack_inputs=unpack_inputs,
             )
             # calculate self_tracin_scores
             self_tracin_scores = tracin.self_influence(

tests/influence/_core/test_tracin_regression.py

@@ -178,8 +178,7 @@ def test_tracin_show_progress(
                 elif mode == "influence":
 
                     tracin.influence(
-                        test_samples,
-                        test_labels,
+                        (test_samples, test_labels),
                         k=None,
                         show_progress=True,
                     )
@@ -196,8 +195,7 @@ def test_tracin_show_progress(
                 elif mode == "k-most":
 
                     tracin.influence(
-                        test_samples,
-                        test_labels,
+                        (test_samples, test_labels),
                         k=2,
                         proponents=True,
                         show_progress=True,
@@ -218,8 +216,7 @@ def test_tracin_show_progress(
                     mock_stderr.truncate(0)
 
                     tracin.influence(
-                        test_samples,
-                        test_labels,
+                        (test_samples, test_labels),
                         k=2,
                         proponents=False,
                         show_progress=True,