def _attribute(
self,
inputs: Tuple[Tensor, ...],
neuron_selector: Union[int, Tuple[int, ...], Callable],
baselines: Tuple[Union[Tensor, int, float], ...],
target: TargetType = None,
additional_forward_args: Any = None,
n_steps: int = 50,
method: str = "riemann_trapezoid",
attribute_to_neuron_input: bool = False,
step_sizes_and_alphas: Union[None, Tuple[List[float],
List[float]]] = None,
) -> Tuple[Tensor, ...]:
num_examples = inputs[0].shape[0]
total_batch = num_examples * n_steps
if step_sizes_and_alphas is None:
# retrieve step size and scaling factor for specified approximation method
step_sizes_func, alphas_func = approximation_parameters(method)
step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
else:
step_sizes, alphas = step_sizes_and_alphas
# Compute scaled inputs from baseline to final input.
scaled_features_tpl = tuple(
torch.cat(
[baseline + alpha * (input - baseline) for alpha in alphas],
dim=0).requires_grad_()
for input, baseline in zip(inputs, baselines))
additional_forward_args = _format_additional_forward_args(
additional_forward_args)
# apply number of steps to additional forward args
# currently, number of steps is applied only to additional forward arguments
# that are nd-tensors. It is assumed that the first dimension is
# the number of batches.
# dim -> (#examples * #steps x additional_forward_args[0].shape[1:], ...)
input_additional_args = (_expand_additional_forward_args(
additional_forward_args, n_steps) if additional_forward_args
is not None else None)
expanded_target = _expand_target(target, n_steps)
# Conductance Gradients - Returns gradient of output with respect to
# hidden layer and hidden layer evaluated at each input.
layer_gradients, layer_eval, input_grads = compute_layer_gradients_and_eval(
forward_fn=self.forward_func,
layer=self.layer,
inputs=scaled_features_tpl,
target_ind=expanded_target,
additional_forward_args=input_additional_args,
gradient_neuron_selector=neuron_selector,
device_ids=self.device_ids,
attribute_to_layer_input=attribute_to_neuron_input,
)
mid_grads = _verify_select_neuron(layer_gradients, neuron_selector)
scaled_input_gradients = tuple(
input_grad * mid_grads.reshape((total_batch, ) + (1, ) *
(len(input_grad.shape) - 1))
for input_grad in input_grads)
# Mutliplies by appropriate step size.
scaled_grads = tuple(
scaled_input_gradient.contiguous().view(n_steps, -1) *
torch.tensor(step_sizes).view(n_steps, 1).to(
scaled_input_gradient.device)
for scaled_input_gradient in scaled_input_gradients)
# Aggregates across all steps for each tensor in the input tuple
total_grads = tuple(
_reshape_and_sum(scaled_grad, n_steps, num_examples,
input_grad.shape[1:])
for (scaled_grad, input_grad) in zip(scaled_grads, input_grads))
if self.multiplies_by_inputs:
# computes attribution for each tensor in input tuple
# attributions has the same dimensionality as inputs
attributions = tuple(total_grad * (input - baseline)
for total_grad, input, baseline in zip(
total_grads, inputs, baselines))
else:
attributions = total_grads
return attributions
def _batch_attribution(
attr_method,
num_examples,
internal_batch_size,
n_steps,
include_endpoint=False,
**kwargs,
):
"""
This method applies internal batching to given attribution method, dividing
the total steps into batches and running each independently and sequentially,
adding each result to compute the total attribution.
Step sizes and alphas are spliced for each batch and passed explicitly for each
call to _attribute.
kwargs include all argument necessary to pass to each attribute call, except
for n_steps, which is computed based on the number of steps for the batch.
include_endpoint ensures that one step overlaps between each batch, which
is necessary for some methods, particularly LayerConductance.
"""
if internal_batch_size < num_examples:
warnings.warn(
"Internal batch size cannot be less than the number of input examples. "
"Defaulting to internal batch size of %d equal to the number of examples."
% num_examples
)
# Number of steps for each batch
step_count = max(1, internal_batch_size // num_examples)
if include_endpoint:
if step_count < 2:
step_count = 2
warnings.warn(
"This method computes finite differences between evaluations at "
"consecutive steps, so internal batch size must be at least twice "
"the number of examples. Defaulting to internal batch size of %d"
" equal to twice the number of examples." % (2 * num_examples)
)
total_attr = None
cumulative_steps = 0
step_sizes_func, alphas_func = approximation_parameters(kwargs["method"])
full_step_sizes = step_sizes_func(n_steps)
full_alphas = alphas_func(n_steps)
while cumulative_steps < n_steps:
start_step = cumulative_steps
end_step = min(start_step + step_count, n_steps)
batch_steps = end_step - start_step
if include_endpoint:
batch_steps -= 1
step_sizes = full_step_sizes[start_step:end_step]
alphas = full_alphas[start_step:end_step]
current_attr = attr_method._attribute(
**kwargs, n_steps=batch_steps, step_sizes_and_alphas=(step_sizes, alphas)
)
if total_attr is None:
total_attr = current_attr
else:
if isinstance(total_attr, Tensor):
total_attr = total_attr + current_attr.detach()
else:
total_attr = tuple(
current.detach() + prev_total
for current, prev_total in zip(current_attr, total_attr)
)
if include_endpoint and end_step < n_steps:
cumulative_steps = end_step - 1
else:
cumulative_steps = end_step
return total_attr
def _attribute(
self,
inputs: Tuple[Tensor, ...],
baselines: Tuple[Union[Tensor, int, float], ...],
target: TargetType = None,
additional_forward_args: Any = None,
n_steps: int = 50,
method: str = "gausslegendre",
attribute_to_layer_input: bool = False,
step_sizes_and_alphas: Union[None, Tuple[List[float],
List[float]]] = None,
) -> Union[Tensor, Tuple[Tensor, ...]]:
if step_sizes_and_alphas is None:
# retrieve step size and scaling factor for specified approximation method
step_sizes_func, alphas_func = approximation_parameters(method)
step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
else:
step_sizes, alphas = step_sizes_and_alphas
# Compute scaled inputs from baseline to final input.
scaled_features_tpl = tuple(
torch.cat(
[baseline + alpha * (input - baseline) for alpha in alphas],
dim=0).requires_grad_()
for input, baseline in zip(inputs, baselines))
additional_forward_args = _format_additional_forward_args(
additional_forward_args)
# apply number of steps to additional forward args
# currently, number of steps is applied only to additional forward arguments
# that are nd-tensors. It is assumed that the first dimension is
# the number of batches.
# dim -> (bsz * #steps x additional_forward_args[0].shape[1:], ...)
input_additional_args = (_expand_additional_forward_args(
additional_forward_args, n_steps) if additional_forward_args
is not None else None)
expanded_target = _expand_target(target, n_steps)
# Returns gradient of output with respect to hidden layer.
layer_gradients, _ = compute_layer_gradients_and_eval(
forward_fn=self.forward_func,
layer=self.layer,
inputs=scaled_features_tpl,
target_ind=expanded_target,
additional_forward_args=input_additional_args,
device_ids=self.device_ids,
attribute_to_layer_input=attribute_to_layer_input,
)
# flattening grads so that we can multiply it with step-size
# calling contiguous to avoid `memory whole` problems
scaled_grads = tuple(
layer_grad.contiguous().view(n_steps, -1) *
torch.tensor(step_sizes).view(n_steps, 1).to(layer_grad.device)
for layer_grad in layer_gradients)
# aggregates across all steps for each tensor in the input tuple
attrs = tuple(
_reshape_and_sum(scaled_grad, n_steps, inputs[0].shape[0],
layer_grad.shape[1:])
for scaled_grad, layer_grad in zip(scaled_grads, layer_gradients))
return _format_output(len(attrs) > 1, attrs)
def attribute(
self,
inputs,
baselines=None,
target=None,
n_steps=500,
method="riemann_trapezoid",
):
r"""
Computes conductance using gradients along the path, applying
riemann's method or gauss-legendre.
The details of the approach can be found here:
https://arxiv.org/abs/1805.12233
Args
inputs: A single high dimensional input tensor, in which
dimension 0 corresponds to number of examples.
baselines: A single high dimensional baseline tensor,
which has the same shape as the input
target: Predicted class index. This is necessary only for
classification use cases
n_steps: The number of steps used by the approximation method
method: Method for integral approximation, one of `riemann_right`,
`riemann_middle`, `riemann_trapezoid` or `gausslegendre`
Return
attributions: Total conductance with respect to each neuron in
output of given layer
"""
if baselines is None:
baselines = 0
gradient_mask = apply_gradient_requirements((inputs, ))
# retrieve step size and scaling factor for specified approximation method
step_sizes_func, alphas_func = approximation_parameters(method)
step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
# compute scaled inputs from baseline to final input.
scaled_features = torch.cat(
[baselines + alpha * (inputs - baselines) for alpha in alphas],
dim=0)
# Conductance Gradients - Returns gradient of output with respect to
# hidden layer, gradient of hidden layer with respect to input,
# and number of hidden units.
input_gradients, mid_layer_gradients, hidden_units = self._conductance_grads(
self.forward_func, scaled_features, target)
# Multiply gradient of hidden layer with respect to input by input - baseline
scaled_input_gradients = torch.repeat_interleave(inputs - baselines,
hidden_units,
dim=0)
scaled_input_gradients = input_gradients * scaled_input_gradients.repeat(
*([len(alphas)] + [1] * (len(scaled_input_gradients.shape) - 1)))
# Sum gradients for each input neuron in order to have total
# for each hidden unit and reshape to match hidden layer shape
summed_input_grads = torch.sum(
scaled_input_gradients,
tuple(range(1, len(
scaled_input_gradients.shape)))).view_as(mid_layer_gradients)
# Rescale gradients of hidden layer by by step size.
scaled_grads = mid_layer_gradients.contiguous().view(
n_steps, -1) * torch.tensor(step_sizes).view(n_steps, 1).to(
mid_layer_gradients.device)
undo_gradient_requirements((inputs, ), gradient_mask)
# Element-wise mutliply gradient of output with respect to hidden layer
# and summed gradients with respect to input (chain rule) and sum across
# stepped inputs.
return _reshape_and_sum(
scaled_grads.view(mid_layer_gradients.shape) * summed_input_grads,
n_steps,
inputs.shape[0],
mid_layer_gradients.shape[1:],
)
def _attribute(
self,
inputs: Tuple[Tensor, ...],
baselines: Tuple[Union[Tensor, int, float], ...],
target: TargetType = None,
additional_forward_args: Any = None,
n_steps: int = 50,
method: str = "gausslegendre",
step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
) -> Tuple[Tensor, ...]:
if step_sizes_and_alphas is None:
# retrieve step size and scaling factor for specified
# approximation method
step_sizes_func, alphas_func = approximation_parameters(method)
step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)
else:
step_sizes, alphas = step_sizes_and_alphas
# scale features and compute gradients. (batch size is abbreviated as bsz)
# scaled_features' dim -> (bsz * #steps x inputs[0].shape[1:], ...)
scaled_features_tpl = tuple(
torch.cat(
[baseline + alpha * (input - baseline) for alpha in alphas], dim=0
).requires_grad_()
for input, baseline in zip(inputs, baselines)
)
additional_forward_args = _format_additional_forward_args(
additional_forward_args
)
# apply number of steps to additional forward args
# currently, number of steps is applied only to additional forward arguments
# that are nd-tensors. It is assumed that the first dimension is
# the number of batches.
# dim -> (bsz * #steps x additional_forward_args[0].shape[1:], ...)
input_additional_args = (
_expand_additional_forward_args(additional_forward_args, n_steps)
if additional_forward_args is not None
else None
)
expanded_target = _expand_target(target, n_steps)
# grads: dim -> (bsz * #steps x inputs[0].shape[1:], ...)
grads = self.gradient_func(
forward_fn=self.forward_func,
inputs=scaled_features_tpl,
target_ind=expanded_target,
additional_forward_args=input_additional_args,
)
# flattening grads so that we can multilpy it with step-size
# calling contiguous to avoid `memory whole` problems
scaled_grads = [
grad.contiguous().view(n_steps, -1)
* torch.tensor(step_sizes).view(n_steps, 1).to(grad.device)
for grad in grads
]
# aggregates across all steps for each tensor in the input tuple
# total_grads has the same dimensionality as inputs
total_grads = tuple(
_reshape_and_sum(
scaled_grad, n_steps, grad.shape[0] // n_steps, grad.shape[1:]
)
for (scaled_grad, grad) in zip(scaled_grads, grads)
)
# computes attribution for each tensor in input tuple
# attributions has the same dimensionality as inputs
if not self.multiplies_by_inputs:
attributions = total_grads
else:
attributions = tuple(
total_grad * (input - baseline)
for total_grad, input, baseline in zip(total_grads, inputs, baselines)
)
return attributions
def _attribute(
self,
inputs: Tuple[Tensor, ...],
baselines: Tuple[Union[Tensor, int, float], ...],
target: TargetType = None,
additional_forward_args: Any = None,
n_steps: int = 50,
method: str = "gausslegendre",
attribute_to_layer_input: bool = False,
step_sizes_and_alphas: Union[None, Tuple[List[float], List[float]]] = None,
) -> Union[Tensor, Tuple[Tensor, ...]]:
num_examples = inputs[0].shape[0]
if step_sizes_and_alphas is None:
# Retrieve scaling factors for specified approximation method
step_sizes_func, alphas_func = approximation_parameters(method)
alphas = alphas_func(n_steps + 1)
else:
_, alphas = step_sizes_and_alphas
# Compute scaled inputs from baseline to final input.
scaled_features_tpl = tuple(
torch.cat(
[baseline + alpha * (input - baseline) for alpha in alphas], dim=0
).requires_grad_()
for input, baseline in zip(inputs, baselines)
)
additional_forward_args = _format_additional_forward_args(
additional_forward_args
)
# apply number of steps to additional forward args
# currently, number of steps is applied only to additional forward arguments
# that are nd-tensors. It is assumed that the first dimension is
# the number of batches.
# dim -> (#examples * #steps x additional_forward_args[0].shape[1:], ...)
input_additional_args = (
_expand_additional_forward_args(additional_forward_args, n_steps + 1)
if additional_forward_args is not None
else None
)
expanded_target = _expand_target(target, n_steps + 1)
# Conductance Gradients - Returns gradient of output with respect to
# hidden layer and hidden layer evaluated at each input.
(layer_gradients, layer_evals,) = compute_layer_gradients_and_eval(
forward_fn=self.forward_func,
layer=self.layer,
inputs=scaled_features_tpl,
additional_forward_args=input_additional_args,
target_ind=expanded_target,
device_ids=self.device_ids,
attribute_to_layer_input=attribute_to_layer_input,
)
# Compute differences between consecutive evaluations of layer_eval.
# This approximates the total input gradient of each step multiplied
# by the step size.
grad_diffs = tuple(
layer_eval[num_examples:] - layer_eval[:-num_examples]
for layer_eval in layer_evals
)
# Element-wise multiply gradient of output with respect to hidden layer
# and summed gradients with respect to input (chain rule) and sum
# across stepped inputs.
attributions = tuple(
_reshape_and_sum(
grad_diff * layer_gradient[:-num_examples],
n_steps,
num_examples,
layer_eval.shape[1:],
)
for layer_gradient, layer_eval, grad_diff in zip(
layer_gradients, layer_evals, grad_diffs
)
)
return _format_output(len(attributions) > 1, attributions)