update tracin influence API

Summary:
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: ee06ef12c79c7bc143b786ae2c5c15d5c8c11684

Author: Fulton Wang

captum/influence/_core/tracincp.py

@@ -187,7 +187,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,
@@ -202,13 +201,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
@@ -270,7 +268,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

@@ -11,6 +11,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,
@@ -162,25 +163,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(

captum/influence/_utils/common.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

@@ -165,19 +165,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)
@@ -208,17 +208,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
@@ -270,7 +270,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:
@@ -364,9 +364,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
@@ -397,9 +397,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

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,

tests/influence/_core/test_tracin_self_influence.py

@@ -258,7 +258,7 @@ def test_tracin_xor(
                     batch_size,
                     criterion,
                 )
-                test_scores = tracin.influence(testset, testlabels)
+                test_scores = tracin.influence((testset, testlabels))
                 idx = torch.argsort(test_scores, dim=1, descending=True)
                 # check that top 5 influences have matching binary classification
                 for i in range(len(idx)):
@@ -288,9 +288,9 @@ def test_tracin_xor(
                     criterion,
                     sample_wise_grads_per_batch=True,
                 )
-                test_scores = tracin.influence(testset, testlabels)
+                test_scores = tracin.influence((testset, testlabels))
                 test_scores_sample_wise_trick = tracin_sample_wise_trick.influence(
-                    testset, testlabels
+                    (testset, testlabels)
                 )
                 assertTensorAlmostEqual(
                     self, test_scores, test_scores_sample_wise_trick

tests/influence/_core/test_tracin_show_progress.py

@@ -2,7 +2,7 @@
 import os
 import unittest
 from functools import partial
-from typing import Callable, Iterator, List, Optional, Union
+from typing import Callable, Iterator, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -14,6 +14,7 @@
 )
 from parameterized import parameterized
 from parameterized.parameterized import param
+from torch import Tensor
 from torch.nn import Module
 from torch.utils.data import DataLoader, Dataset
 
@@ -344,3 +345,12 @@ def build_test_name_func(args_to_skip: Optional[List[str]] = None):
     """
 
     return partial(generate_test_name, args_to_skip=args_to_skip)
+
+
+def _format_batch_into_tuple(
+    inputs: Union[Tuple, Tensor], targets: Tensor, unpack_inputs: bool
+):
+    if unpack_inputs:
+        return (*inputs, targets)
+    else:
+        return (inputs, targets)