erge remote-tracking branch 'upstream/master' into insights-occlusion

Author: edward-io

captum/attr/_core/deep_lift.py

@@ -540,7 +540,9 @@ def pre_hook(module: Module, baseline_inputs_add_args: Tuple) -> Tuple:
         def forward_hook(module: Module, inputs: Tuple, outputs: Tensor):
             return torch.stack(torch.chunk(outputs, 2), dim=1)
 
-        if isinstance(self.model, nn.DataParallel):
+        if isinstance(
+            self.model, (nn.DataParallel, nn.parallel.DistributedDataParallel)
+        ):
             return [
                 self.model.module.register_forward_pre_hook(pre_hook),  # type: ignore
                 self.model.module.register_forward_hook(forward_hook),

captum/insights/config.py

@@ -73,16 +73,18 @@ def _str_to_tuple(s):
         post_process={"n_steps": int},
     ),
     FeatureAblation.get_name(): ConfigParameters(
-        params={"perturbations_per_eval": NumberConfig(value=1, limit=(1, 100))}
+        params={"perturbations_per_eval": NumberConfig(value=1, limit=(1, 100))},
     ),
     Occlusion.get_name(): ConfigParameters(
         params={
             "sliding_window_shapes": StrConfig(value=""),
             "strides": StrConfig(value=""),
+            "perturbations_per_eval": NumberConfig(value=1, limit=(1, 100)),
         },
         post_process={
             "sliding_window_shapes": _str_to_tuple,
             "strides": _str_to_tuple,
+            "perturbations_per_eval": int,
         },
     ),
 }

captum/metrics/__init__.py

@@ -1,3 +1,3 @@
 #!/usr/bin/env python3
 
-from ._core.infidelity import infidelity  # noqa
+from ._core.infidelity import infidelity, infidelity_perturb_func_decorator  # noqa

captum/metrics/_core/infidelity.py

@@ -16,6 +16,79 @@
 from .._utils.batching import _divide_and_aggregate_metrics
 
 
+def infidelity_perturb_func_decorator(pertub_func):
+    r"""
+    An auxiliary, decorator function that helps with computing
+    perturbations given perturbed inputs. It can be useful for cases
+    when `pertub_func` returns only perturbed inputs and we
+    internally compute the perturbations as
+    (input - perturbed_input) / (input - baseline).
+
+    If users decorate their `pertub_func` with
+    `@infidelity_perturb_func_decorator` function then their `pertub_func`
+    needs to only return perturbed inputs.
+
+    Note that if your attribution algorithm is inherently local such as
+    Saliency maps you should not use the decorator because the decorator
+    always divides by (input - baseline) and that is unnecessary for local
+    methods.
+    Args:
+
+        pertub_func(callable): Input perturbation function that takes inputs
+            and optionally baselines and returns perturbed inputs
+
+    Returns:
+
+        default_perturb_func(callable): Internal default perturbation
+        function that computes the perturbations internally and returns
+        perturbations and perturbed inputs.
+
+    Examples::
+        >>> @infidelity_perturb_func_decorator
+        >>> def perturb_fn(inputs):
+        >>>    noise = torch.tensor(np.random.normal(0, 0.003, inputs.shape)).float()
+        >>>    return inputs - noise
+        >>> # Computes infidelity score using `perturb_fn`
+        >>> infidelity = infidelity_attr(model, perturb_fn, input, ...)
+
+    """
+
+    def default_perturb_func(inputs, baselines=None):
+        r"""
+        """
+        inputs_perturbed = (
+            pertub_func(inputs, baselines)
+            if baselines is not None
+            else pertub_func(inputs)
+        )
+        inputs_perturbed = _format_tensor_into_tuples(inputs_perturbed)
+        inputs = _format_tensor_into_tuples(inputs)
+        baselines = _format_tensor_into_tuples(baselines)
+        if baselines is None:
+            perturbations = tuple(
+                safe_div(
+                    input - input_perturbed,
+                    input,
+                    torch.tensor(1.0, device=input.device),
+                )
+                for input, input_perturbed in zip(inputs, inputs_perturbed)
+            )
+        else:
+            perturbations = tuple(
+                safe_div(
+                    input - input_perturbed,
+                    input - baseline,
+                    torch.tensor(1.0, device=input.device),
+                )
+                for input, input_perturbed, baseline in zip(
+                    inputs, inputs_perturbed, baselines
+                )
+            )
+        return perturbations, inputs_perturbed
+
+    return default_perturb_func
+
+
 def infidelity(
     forward_func,
     perturb_func,
@@ -26,7 +99,6 @@ def infidelity(
     target=None,
     n_samples=10,
     max_examples_per_batch=None,
-    perturb_func_custom=False,
 ):
     r"""
     Explanation infidelity represents the expected mean-squared error
@@ -62,35 +134,52 @@ def infidelity(
                 The perturbation function of model inputs. This function takes
                 model inputs and optionally baselines as input arguments and returns
                 either a tuple of perturbations and perturbed inputs or just
-                perturbed inputs. If `perturb_func` returns only perturbed inputs
-                then the users have to set the `perturb_func_custom=True`, this
-                will allow us to compute the perturbations internally both for local
-                and global infidelity and makes sense
-                to use only if input attributions are global attributions.
+                perturbed inputs. For example:
+
+                 def my_perturb_func(inputs):
+                    <MY-LOGIC-HERE>
+                    return perturbations, perturbed_inputs
+
+                If we want to only return perturbed inputs and compute
+                perturbations internally then we can wrap perturb_func with
+                `infidelity_perturb_func_decorator` decorator such as:
+
+                from captum.metrics import infidelity_perturb_func_decorator
+                @infidelity_perturb_func_decorator
+                def my_perturb_func(inputs):
+                    <MY-LOGIC-HERE>
+                    return perturbed_inputs
+
+                In this case we compute perturbations by dividing
+                (input - perturbed_input) by (input - baselines) and the user needs to
+                only return perturbed inputs in `perturb_func` as described above.
+
+                `infidelity_perturb_func_decorator` makes sense to use only for global
+                attribution algorithms such as integrated gradients, deeplift, etc.
+                In case user has a local attribution algorithm or decides to compute
+                perturbations and perturbed inputs in `perturb_func` then they must not
+                use `infidelity_perturb_func_decorator`.
 
                 If there are more than one inputs passed to infidelity function those
                 will be passed to `perturb_func` as tuples in the same order as they
                 are passed to infidelity function.
 
-                In case `perturb_func_custom=False` and if inputs
+                If inputs
                  - is a single tensor, the function needs to return a tuple
                     of perturbations and perturbed input such as:
                       perturb, perturbed_input
+
+                      and only perturbed_input in case
+                      `infidelity_perturb_func_decorator`
+                      is used.
                  - is a tuple of tensors,
                     corresponding perturbations and perturbed inputs must be computed
                     and returned as tuples in the following format:
                         (perturb1, perturb2, ... perturbN), (perturbed_input1,
                          perturbed_input2, ... perturbed_inputN)
-
-                In case `perturb_func_custom=True` and if inputs
-                 - is a single tensor, the function needs to return
-                     only perturbed input
-                       perturbed_input
-                 - is a tuple of tensors,
-                    corresponding perturbed inputs must be computed and
-                    returned as tuples in the following format:
-                       (perturbed_input1, perturbed_input2, ... perturbed_inputN)
-
+                    Similar to previous case here as well we need to return only
+                    perturbed inputs in case `infidelity_perturb_func_decorator`
+                    decorates out perturb_func
                 It is important to note that for performance reasons `perturb_func`
                 isn't called for each example individually but on a batch of
                 input examples that are repeated `max_examples_per_batch / batch_size`
@@ -164,9 +253,8 @@ def infidelity(
                 tensor as well. If inputs is provided as a tuple of tensors
                 then attributions will be tuples of tensors as well.
 
-                If `perturb_func_custom=True` then we internally divide global
-                attribution values by (input - baselines) and the user needs to
-                only return perturbed inputs in `perturb_func`.
+                For more details on when to use `infidelity_perturb_func_decorator`,
+                please, read the documentation about `perturb_func`
 
         additional_forward_args (any, optional): If the forward function
                 requires additional arguments other than the inputs for
@@ -220,17 +308,6 @@ def infidelity(
                 examples are processed together.
 
                 Default: None
-        perturb_func_custom (boolean, optional): A flag that indicates whether
-                to use default perturbation logic that always divides the
-                attributions by (input - baseline). If this flag
-                is True then `perturb_func` needs to return only the
-                perturbed inputs.
-                The perturbations will be computed internally by
-                `default_perturb_func`. This makes sense to use only with
-                global attribution values because otherwise there is no need
-                to divide the attributions by (input - baseline).
-
-                Default: False
     Returns:
 
         infidelities (tensor): A tensor of scalar infidelity scores per
@@ -254,31 +331,6 @@ def infidelity(
         >>> infidelity = infidelity_attr(net, perturb_fn, input, attribution)
     """
 
-    def default_perturb_func(inputs, inputs_perturbed, baselines=None):
-        r"""
-        """
-        if baselines is None:
-            perturbations = tuple(
-                safe_div(
-                    input - input_perturbed,
-                    input,
-                    torch.tensor(1.0, device=input.device),
-                )
-                for input, input_perturbed in zip(inputs, inputs_perturbed)
-            )
-        else:
-            perturbations = tuple(
-                safe_div(
-                    input - input_perturbed,
-                    input - baseline,
-                    torch.tensor(1.0, device=input.device),
-                )
-                for input, input_perturbed, baseline in zip(
-                    inputs, inputs_perturbed, baselines
-                )
-            )
-        return perturbations, inputs_perturbed
-
     def _generate_perturbations(current_n_samples):
         r"""
         The perturbations are generated for each example `current_n_samples` times.
@@ -308,6 +360,7 @@ def call_perturb_func():
         inputs_expanded = tuple(
             torch.repeat_interleave(input, current_n_samples, dim=0) for input in inputs
         )
+
         if baselines is not None:
             baselines_expanded = tuple(
                 baseline.repeat_interleave(current_n_samples, dim=0)
@@ -320,16 +373,7 @@ def call_perturb_func():
         else:
             baselines_expanded = None
 
-        perturb_func_out = call_perturb_func()
-
-        if perturb_func_custom:
-            return default_perturb_func(
-                inputs_expanded,
-                _format_tensor_into_tuples(perturb_func_out),
-                baselines=baselines_expanded,
-            )
-        else:
-            return perturb_func_out
+        return call_perturb_func()
 
     def _validate_inputs_and_perturbations(inputs, inputs_perturbed, perturbations):
         # asserts the sizes of the perturbations and inputs

docs/algorithms_comparison_matrix.md

@@ -213,3 +213,5 @@ Please, scroll to the right for more details.
 **^ Including Layer Variant**
 
 **˚ Including Neuron Variant**
+
+<a href="/img/algorithms_comparison_matrix.png">Algorithm Comparison Matrix.png</a>

scripts/update_versions_html.py

@@ -35,8 +35,8 @@ def prepend_url(a_tag, base_url, version):
 
         # nav
         nav_links = soup.find("nav").findAll("a")
-        for l in nav_links:
-            l.attrs["href"] = prepend_url(l, base_url, v)
+        for link in nav_links:
+            link.attrs["href"] = prepend_url(link, base_url, v)
 
         # version link
         t = soup.find("h2", {"class": "headerTitleWithLogo"}).find_next("a")

tests/attr/test_data_parallel.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 import copy
+import os
 from enum import Enum
 from typing import Any, Callable, Dict, Optional, Tuple, Type, cast
 
 import torch
+import torch.distributed as dist
 from torch import Tensor
 from torch.nn import Module
 
@@ -22,14 +24,17 @@
 from .helpers.gen_test_utils import gen_test_name, parse_test_config
 from .helpers.test_config import config
 
-
 """
 Tests in this file are dynamically generated based on the config
 defined in tests/attr/helpers/test_config.py. To add new test cases,
 read the documentation in test_config.py and add cases based on the
 schema described there.
 """
 
+# Distributed Data Parallel env setup
+os.environ["MASTER_ADDR"] = "127.0.0.1"
+os.environ["MASTER_PORT"] = "29500"
+
 
 class DataParallelCompareMode(Enum):
     """
@@ -44,6 +49,7 @@ class DataParallelCompareMode(Enum):
     cpu_cuda = 1
     data_parallel_default = 2
     data_parallel_alt_dev_ids = 3
+    dist_data_parallel = 4
 
 
 class DataParallelMeta(type):
@@ -146,6 +152,13 @@ def data_parallel_test_assert(self) -> None:
                     ),
                 )
                 args_1, args_2 = cuda_args, cuda_args
+            elif mode is DataParallelCompareMode.dist_data_parallel:
+                dist.init_process_group(backend="gloo", rank=0, world_size=1)
+                model_1, model_2 = (
+                    cuda_model,
+                    torch.nn.parallel.DistributedDataParallel(cuda_model),
+                )
+                args_1, args_2 = cuda_args, cuda_args
             else:
                 raise AssertionError("DataParallel compare mode type is not valid.")
 
@@ -219,8 +232,11 @@ def data_parallel_test_assert(self) -> None:
                     self, attributions_1, attributions_2, mode="max", delta=dp_delta
                 )
 
+            if mode is DataParallelCompareMode.dist_data_parallel:
+                dist.destroy_process_group()
+
         return data_parallel_test_assert
 
 
-class DataParallelTest(BaseGPUTest):
+class DataParallelTest(BaseGPUTest, metaclass=DataParallelMeta):
     pass

tests/attr/test_hook_removal.py

@@ -52,7 +52,7 @@ class HookRemovalMeta(type):
     """
 
     def __new__(cls, name: str, bases: Tuple, attrs: Dict):
-        created_tests = {}
+        created_tests: Dict[Tuple[Type[Attribution], HookRemovalMode], bool] = {}
         for test_config in config:
             (algorithms, model, args, layer, noise_tunnel, _,) = parse_test_config(
                 test_config