-
Notifications
You must be signed in to change notification settings - Fork 1
/
integrated_gradients.py
82 lines (72 loc) · 3.12 KB
/
integrated_gradients.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from typing import Callable, Optional, Tuple, Union, Sequence, Dict
from torch import Tensor
from torch.nn.modules import Module
from captum.attr import IntegratedGradients as CaptumIntegratedGradients
from captum.attr import LayerIntegratedGradients as CaptumLayerIntegratedGradients
from pnpxai.core.detector.types import Linear, Convolution, Attention
from pnpxai.utils import format_into_tuple, format_out_tuple_if_single
from pnpxai.explainers.utils.baselines import BaselineMethodOrFunction, BaselineFunction
from .base import Explainer
from .utils import captum_wrap_model_input
class IntegratedGradients(Explainer):
SUPPORTED_MODULES = [Linear, Convolution, Attention]
def __init__(
self,
model: Module,
n_steps: int = 20,
baseline_fn: Union[BaselineMethodOrFunction,
Tuple[BaselineMethodOrFunction]] = 'zeros',
layer: Optional[Callable[[Tuple[Tensor]],
Union[Tensor, Tuple[Tensor]]]] = None,
forward_arg_extractor: Optional[Callable[[
Tuple[Tensor]], Union[Tensor, Tuple[Tensor]]]] = None,
additional_forward_arg_extractor: Optional[Callable[[
Tuple[Tensor]], Union[Tensor, Tuple[Tensor]]]] = None,
) -> None:
super().__init__(model, forward_arg_extractor, additional_forward_arg_extractor)
self.layer = layer
self.n_steps = n_steps
self.baseline_fn = baseline_fn
@property
def _layer_explainer(self) -> CaptumLayerIntegratedGradients:
wrapped_model = captum_wrap_model_input(self.model)
layers = [
wrapped_model.input_maps[layer] if isinstance(layer, str)
else layer for layer in self.layer
] if isinstance(self.layer, Sequence) else self.layer
return CaptumLayerIntegratedGradients(
forward_func=wrapped_model,
layer=layers,
)
@property
def _explainer(self) -> CaptumIntegratedGradients:
return CaptumIntegratedGradients(forward_func=self.model)
@property
def explainer(self) -> Union[CaptumIntegratedGradients, CaptumLayerIntegratedGradients]:
if self.layer is None:
return self._explainer
return self._layer_explainer
def attribute(
self,
inputs: Union[Tensor, Tuple[Tensor]],
targets: Tensor
) -> Union[Tensor, Tuple[Tensor]]:
forward_args, additional_forward_args = self._extract_forward_args(
inputs)
forward_args = format_into_tuple(forward_args)
baselines = format_into_tuple(self._get_baselines(forward_args))
attrs = self.explainer.attribute(
inputs=forward_args,
baselines=baselines,
target=targets,
additional_forward_args=additional_forward_args,
n_steps=self.n_steps,
)
if isinstance(attrs, tuple):
attrs = format_out_tuple_if_single(attrs)
return attrs
def get_tunables(self) -> Dict[str, Tuple[type, dict]]:
return {
'n_steps': (int, {'low': 10, 'high': 100, 'step': 10}),
'baseline_fn': (BaselineFunction, {}),
}