Optimization in SVS when formatting total_attrib (#1648)

Summary:
Pull Request resolved: #1648

This diff focuses on optimizing the Shapley Value Sampling by formatting a default dictionary to list all of the feature indices with a curr_mask that is non-zero, and only updates the total_attrib for values which should be updated rather than iterating through the entire feature space which will mostly be unaffected.

Reviewed By: vivekmig

Differential Revision: D81800703

fbshipit-source-id: 53b1a5f0589b4bf6392139d4772e5ec68887e4ae

Author: jjuncho

captum/attr/_core/shapley_value.py

@@ -15,6 +15,7 @@
     List,
     Optional,
     Sequence,
+    Set,
     Tuple,
     Union,
 )
@@ -382,6 +383,7 @@ def attribute(
                     current_add_args,
                     current_target,
                     current_masks,
+                    current_feat_list,
                 ) in self._perturbation_generator(
                     inputs_tuple,
                     additional_forward_args,
@@ -423,8 +425,7 @@ def attribute(
                             all_eval[num_examples:] - all_eval[:-num_examples]
                         ).to(inputs_tuple[0].device)
                         prev_results = all_eval[-num_examples:]
-
-                    for j in range(len(total_attrib)):
+                    for j in current_feat_list:
                         # format eval_diff to shape
                         # (n_perturb, *output_shape, 1,.. 1)
                         # where n_perturb may not be perturb_per_eval
@@ -545,6 +546,7 @@ def attribute_future(
                     current_add_args,
                     current_target,
                     current_masks,
+                    _,
                 ) in self._perturbation_generator(
                     inputs_tuple,
                     additional_forward_args,
@@ -827,7 +829,9 @@ def _perturbation_generator(
         input_masks: Tuple[Tensor, ...],
         feature_permutation: Sequence[int],
         perturbations_per_eval: int,
-    ) -> Iterable[Tuple[Tuple[Tensor, ...], object, TargetType, Tuple[Tensor, ...]]]:
+    ) -> Iterable[
+        Tuple[Tuple[Tensor, ...], object, TargetType, Tuple[Tensor, ...], Set[int]]
+    ]:
         """
         This method is a generator which yields each perturbation to be evaluated
         including inputs, additional_forward_args, targets, and mask.
@@ -837,6 +841,7 @@ def _perturbation_generator(
         current_tensors = baselines
         current_tensors_list = []
         current_mask_list = []
+        current_feat_list = set()
 
         # Compute repeated additional args and targets
         additional_args_repeated = (
@@ -869,6 +874,7 @@ def _perturbation_generator(
                     device=inputs[0].device,
                 )
             )
+            current_feat_list.update(feat_tensor_index_map[feature_permutation[i]])
 
             if len(current_tensors_list) == perturbations_per_eval:
                 if len(current_tensors_list) > 1:
@@ -888,9 +894,11 @@ def _perturbation_generator(
                     additional_args_repeated,
                     target_repeated,
                     combined_masks,
+                    current_feat_list,
                 )
                 current_tensors_list = []
                 current_mask_list = []
+                current_feat_list = set()
 
         # Create batch with remaining evaluations, may not be a complete batch
         # (= perturbations_per_eval)
@@ -916,6 +924,7 @@ def _perturbation_generator(
                 additional_args_repeated,
                 target_repeated,
                 combined_masks,
+                current_feat_list,
             )
 
     def _get_n_evaluations(