Expose attribution forward progress for runtime estimation (#1933)

captum/_utils/progress.py

@@ -26,6 +26,25 @@
 
 T = TypeVar("T")
 IterableType = TypeVar("IterableType", covariant=True)
+ProgressCallback = Callable[[Optional[int], float], None]
+
+
+@runtime_checkable
+class ForwardProgressListener(Protocol):
+    """Receives the planned and completed model forwards for an attribution."""
+
+    def on_plan(self, total: int) -> None: ...
+
+    def on_forward_complete(self) -> None: ...
+
+    def on_forward_failed(self) -> None: ...
+
+
+@runtime_checkable
+class SupportsForwardProgressPlan(Protocol):
+    """Attribution algorithm that can predict its model-forward count."""
+
+    def expected_forward_count(self, *args: Any, **kwargs: Any) -> int: ...
 
 
 @runtime_checkable

captum/attr/_core/feature_ablation.py

@@ -210,7 +210,15 @@ def _should_skip_inputs_and_warn(
     Returns:
         bool: True if the feature group should be skipped, False otherwise.
     """
-    should_skip = False
+    should_skip = _should_skip_inputs(
+        current_feature_idxs,
+        feature_idx_to_tensor_idx,
+        formatted_inputs,
+        min_examples_per_batch_grouped,
+    )
+    if not should_skip:
+        return False
+
     all_empty = True
     tensor_idx_list = []
     for feature_idx in current_feature_idxs:
@@ -221,15 +229,12 @@ def _should_skip_inputs_and_warn(
         if min_examples_per_batch_grouped is not None and (
             formatted_inputs[tensor_idx].shape[0] < min_examples_per_batch_grouped
         ):
-            should_skip = True
-            break
-    if should_skip:
-        logger.warning(
-            f"Skipping feature group {current_feature_idxs} since it contains "
-            f"at least one input tensor with 0th dim less than "
-            f"{min_examples_per_batch_grouped}"
-        )
-        return True
+            logger.warning(
+                f"Skipping feature group {current_feature_idxs} since it contains "
+                f"at least one input tensor with 0th dim less than "
+                f"{min_examples_per_batch_grouped}"
+            )
+            return True
     if all_empty:
         logger.info(
             f"Skipping feature group {current_feature_idxs} since all "
@@ -239,6 +244,28 @@ def _should_skip_inputs_and_warn(
     return False
 
 
+def _should_skip_inputs(
+    current_feature_idxs: list[int],
+    feature_idx_to_tensor_idx: dict[int, list[int]],
+    formatted_inputs: tuple[Tensor, ...],
+    min_examples_per_batch_grouped: int | None = None,
+) -> bool:
+    """Return whether attribution will skip this group without logging."""
+    tensor_idxs = {
+        tensor_idx
+        for feature_idx in current_feature_idxs
+        for tensor_idx in feature_idx_to_tensor_idx[feature_idx]
+    }
+    if not tensor_idxs:
+        return True
+    if all(torch.numel(formatted_inputs[index]) == 0 for index in tensor_idxs):
+        return True
+    return min_examples_per_batch_grouped is not None and any(
+        formatted_inputs[index].shape[0] < min_examples_per_batch_grouped
+        for index in tensor_idxs
+    )
+
+
 class FeatureAblation(PerturbationAttribution):
     """
     A perturbation based approach to computing attribution, involving
@@ -291,6 +318,58 @@ def __init__(
         # feature group (not parameterized by `_min_examples_per_batch_grouped`).
         self._min_examples_per_batch_grouped: Optional[int] = None
 
+    def expected_forward_count(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: object | None = None,
+        feature_mask: Tensor | tuple[Tensor, ...] | None = None,
+        perturbations_per_eval: int = 1,
+        show_progress: bool = False,
+        run_forward_on_skip: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        """Return the exact number of model forwards for this attribution."""
+        if type(self) is not FeatureAblation:
+            raise NotImplementedError(
+                f"{type(self).__name__} must provide its own exact forward plan."
+            )
+        del baselines, target, additional_forward_args, show_progress
+        formatted_inputs, _ = _format_input_baseline(inputs, None)
+        formatted_feature_mask = _format_feature_mask(feature_mask, formatted_inputs)
+        return self._expected_forward_count_from_formatted(
+            formatted_inputs,
+            formatted_feature_mask,
+            perturbations_per_eval,
+            run_forward_on_skip,
+            **kwargs,
+        )
+
+    def _expected_forward_count_from_formatted(
+        self,
+        formatted_inputs: tuple[Tensor, ...],
+        formatted_feature_mask: tuple[Tensor, ...],
+        perturbations_per_eval: int,
+        run_forward_on_skip: bool,
+        **kwargs: Any,
+    ) -> int:
+        feature_idx_to_tensor_idx = self._get_feature_idx_to_tensor_idx(
+            formatted_feature_mask, **kwargs
+        )
+        feature_idxs = list(feature_idx_to_tensor_idx)
+        perturbation_forwards = 0
+        for start in range(0, len(feature_idxs), perturbations_per_eval):
+            current_feature_idxs = feature_idxs[start : start + perturbations_per_eval]
+            if run_forward_on_skip or not _should_skip_inputs(
+                current_feature_idxs,
+                feature_idx_to_tensor_idx,
+                formatted_inputs,
+                self._min_examples_per_batch_grouped,
+            ):
+                perturbation_forwards += 1
+        return 1 + perturbation_forwards
+
     @log_usage(part_of_slo=True)
     @torch.no_grad()
     def attribute(

captum/attr/_core/feature_permutation.py

@@ -9,9 +9,14 @@
 from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
 
 import torch
-from captum._utils.common import _format_output, _format_tensor_into_tuples
+from captum._utils.common import (
+    _format_feature_mask,
+    _format_output,
+    _format_tensor_into_tuples,
+)
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
 from captum.attr._core.feature_ablation import FeatureAblation
+from captum.attr._utils.common import _format_input_baseline
 from captum.log import log_usage
 from torch import Tensor
 from torch.futures import Future
@@ -104,6 +109,40 @@ def __init__(
         # `_min_examples_per_batch_grouped`.
         self._min_examples_per_batch_grouped = 2
 
+    def expected_forward_count(
+        self,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: object | None = None,
+        feature_mask: Tensor | tuple[Tensor, ...] | None = None,
+        perturbations_per_eval: int = 1,
+        show_progress: bool = False,
+        run_forward_on_skip: bool = False,
+        n_samples: int = 1,
+        **kwargs: Any,
+    ) -> int:
+        """Return the aggregate forward count across all permutation samples."""
+        if type(self) is not FeaturePermutation:
+            raise NotImplementedError(
+                f"{type(self).__name__} must provide its own exact forward plan."
+            )
+        del baselines
+        assert (
+            isinstance(n_samples, int) and n_samples >= 1
+        ), "n_samples must be an integer and at least 1."
+        del target, additional_forward_args, show_progress
+        formatted_inputs, _ = _format_input_baseline(inputs, None)
+        formatted_feature_mask = _format_feature_mask(feature_mask, formatted_inputs)
+        per_sample = self._expected_forward_count_from_formatted(
+            formatted_inputs,
+            formatted_feature_mask,
+            perturbations_per_eval,
+            run_forward_on_skip,
+            **kwargs,
+        )
+        return n_samples * per_sample
+
     # suppressing error caused by the child class not having a matching
     # signature to the parent
     @log_usage(part_of_slo=True)

tests/attr/test_feature_ablation.py

@@ -77,6 +77,40 @@ def test_simple_ablation(self) -> None:
             ablation_algo, inp, [[80.0, 200.0, 120.0]], perturbations_per_eval=(1, 2, 3)
         )
 
+    def test_forward_plan_matches_batched_sparse_mask_calls(self) -> None:
+        forward_calls = 0
+
+        def counted_forward(inputs: Tensor) -> Tensor:
+            nonlocal forward_calls
+            forward_calls += 1
+            return inputs.sum(dim=1)
+
+        algorithm = FeatureAblation(counted_forward)
+        inputs = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        feature_mask = torch.tensor([[0, 2, 5]])
+        planned = algorithm.expected_forward_count(
+            inputs,
+            feature_mask=feature_mask,
+            perturbations_per_eval=2,
+        )
+
+        algorithm.attribute(
+            inputs,
+            feature_mask=feature_mask,
+            perturbations_per_eval=2,
+        )
+
+        self.assertEqual(planned, 3)
+        self.assertEqual(planned, forward_calls)
+
+    def test_subclass_without_exact_planner_is_rejected(self) -> None:
+        class UnsupportedAblation(FeatureAblation):
+            pass
+
+        algorithm = UnsupportedAblation(lambda inputs: inputs.sum(dim=1))
+        with self.assertRaisesRegex(NotImplementedError, "exact forward plan"):
+            algorithm.expected_forward_count(torch.tensor([[1.0, 2.0]]))
+
     def test_simple_ablation_int_to_int(self) -> None:
         ablation_algo = FeatureAblation(BasicModel())
         inp = torch.tensor([[-3, 1, 2]])

tests/attr/test_feature_permutation.py

@@ -61,6 +61,45 @@ def test_perm_fn_single_feature(self) -> None:
                 self._check_perm_fn_with_mask(inp, flat_mask.view_as(inp[0]))
                 flat_mask[i] = 0
 
+    def test_forward_plan_aggregates_samples_and_matches_calls(self) -> None:
+        forward_calls = 0
+
+        def counted_forward(inputs: Tensor) -> Tensor:
+            nonlocal forward_calls
+            forward_calls += 1
+            return inputs.sum(dim=1)
+
+        algorithm = FeaturePermutation(counted_forward)
+        inputs = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+        feature_mask = torch.tensor([[0, 2, 5]])
+        planned = algorithm.expected_forward_count(
+            inputs,
+            feature_mask=feature_mask,
+            perturbations_per_eval=2,
+            n_samples=3,
+        )
+
+        algorithm.attribute(
+            inputs,
+            feature_mask=feature_mask,
+            perturbations_per_eval=2,
+            n_samples=3,
+        )
+
+        self.assertEqual(planned, 9)
+        self.assertEqual(planned, forward_calls)
+
+    def test_forward_plan_requires_subclass_override(self) -> None:
+        class DerivedFeaturePermutation(FeaturePermutation):
+            pass
+
+        algorithm = DerivedFeaturePermutation(lambda inputs: inputs.sum(dim=1))
+
+        with self.assertRaisesRegex(
+            NotImplementedError, "must provide its own exact forward plan"
+        ):
+            algorithm.expected_forward_count(torch.tensor([[1.0, 2.0]]))
+
     def test_perm_fn_broadcastable_masks(self) -> None:
         batch_size = 5
         inp_size = (3, 20, 30)