def compute_gradients(
forward_fn: Callable,
inputs: Union[Tensor, Tuple[Tensor, ...]],
target_ind: TargetType = None,
additional_forward_args: Any = None,
) -> Tuple[Tensor, ...]:
r"""
Computes gradients of the output with respect to inputs for an
arbitrary forward function.
Args:
forward_fn: forward function. This can be for example model's
forward function.
input: Input at which gradients are evaluated,
will be passed to forward_fn.
target_ind: Index of the target class for which gradients
must be computed (classification only).
additional_forward_args: Additional input arguments that forward
function requires. It takes an empty tuple (no additional
arguments) if no additional arguments are required
"""
with torch.autograd.set_grad_enabled(True):
# runs forward pass
outputs = _run_forward(forward_fn, inputs, target_ind,
additional_forward_args)
assert outputs[0].numel() == 1, (
"Target not provided when necessary, cannot"
" take gradient with respect to multiple outputs.")
# torch.unbind(forward_out) is a list of scalar tensor tuples and
# contains batch_size * #steps elements
grads = torch.autograd.grad(torch.unbind(outputs), inputs)
return grads
def forward_fn():
return _run_forward(forward_func, inputs, target,
additional_forward_args)
def _calculate_vis_output(self,
inputs,
additional_forward_args,
label,
target=None) -> Optional[VisualizationOutput]:
net = self.models[0] # TODO process multiple models
# initialize baselines
baseline_transforms_len = len(self.features[0].baseline_transforms
or [])
baselines = [[None] * len(self.features)
for _ in range(baseline_transforms_len)]
transformed_inputs = list(inputs)
# transformed_inputs = list([i.clone() for i in inputs])
for feature_i, feature in enumerate(self.features):
if feature.input_transforms is not None:
transformed_inputs[feature_i] = self._transform(
feature.input_transforms, transformed_inputs[feature_i],
True)
if feature.baseline_transforms is not None:
assert baseline_transforms_len == len(
feature.baseline_transforms
), "Must have same number of baselines across all features"
for baseline_i, baseline_transform in enumerate(
feature.baseline_transforms):
baselines[baseline_i][feature_i] = self._transform(
baseline_transform, transformed_inputs[feature_i],
True)
outputs = _run_forward(
net,
tuple(transformed_inputs),
additional_forward_args=additional_forward_args,
)
if self.score_func is not None:
outputs = self.score_func(outputs)
if outputs.nelement() == 1:
scores = outputs
predicted = scores.round().to(torch.int)
else:
scores, predicted = outputs.topk(min(4, outputs.shape[-1]))
scores = scores.cpu().squeeze(0)
predicted = predicted.cpu().squeeze(0)
if label is not None and len(label) > 0:
actual_label_output = OutputScore(score=100,
index=label[0],
label=self.classes[label[0]])
else:
actual_label_output = None
predicted_scores = self._get_labels_from_scores(scores, predicted)
# Filter based on UI configuration
if not self._should_keep_prediction(predicted_scores,
actual_label_output):
return None
baselines = [tuple(b) for b in baselines]
if target is None:
target = predicted_scores[0].index if len(
predicted_scores) > 0 else None
# attributions are given per input*
# inputs given to the model are described via `self.features`
#
# *an input contains multiple features that represent it
# e.g. all the pixels that describe an image is an input
attrs_per_input_feature = self._calculate_attribution(
net, baselines, tuple(transformed_inputs), additional_forward_args,
target)
net_contrib = self._calculate_net_contrib(attrs_per_input_feature)
# the features per input given
features_per_input = [
feature.visualize(attr, data, contrib)
for feature, attr, data, contrib in zip(
self.features, attrs_per_input_feature, inputs, net_contrib)
]
return VisualizationOutput(
feature_outputs=features_per_input,
actual=actual_label_output,
predicted=predicted_scores,
active_index=target
if target is not None else actual_label_output.index,
)
def _forward_layer_distributed_eval(
forward_fn: Callable,
inputs: Union[Tensor, Tuple[Tensor, ...]],
layer: Module,
target_ind: TargetType = None,
additional_forward_args: Any = None,
attribute_to_layer_input: bool = False,
forward_hook_with_return: bool = False,
) -> Union[Tuple[Dict[device, Tuple[Tensor, ...]], Tensor, bool], Tuple[Dict[
device, Tuple[Tensor, ...]], bool], ]:
r"""
A helper function that allows to set a hook on model's `layer`, run the forward
pass and returns intermediate layer results, stored in a dictionary,
and optionally also the output of the forward function. The keys in the
dictionary are the device ids and the values are corresponding intermediate layer
results, either the inputs or the outputs of the layer depending on whether we set
`attribute_to_layer_input` to True or False.
This is especially useful when we execute forward pass in a distributed setting,
using `DataParallel`s for example.
"""
saved_layer = {}
is_eval_tuple = False
lock = threading.Lock()
# Set a forward hook on specified module and run forward pass to
# get layer output tensor(s).
# For DataParallel models, each partition adds entry to dictionary
# with key as device and value as corresponding Tensor.
def forward_hook(module, inp, out=None):
nonlocal is_eval_tuple
eval_tsrs = inp if attribute_to_layer_input else out
is_eval_tuple = isinstance(eval_tsrs, tuple)
if not is_eval_tuple:
eval_tsrs = (eval_tsrs, )
with lock:
nonlocal saved_layer
# Note that cloning behaviour of `eval_tsr` is different
# when `forward_hook_with_return` is set to True. This is because
# otherwise `backward()` on the last output layer won't execute.
if forward_hook_with_return:
saved_layer[eval_tsrs[0].device] = eval_tsrs
eval_tsrs_to_return = tuple(eval_tsr.clone()
for eval_tsr in eval_tsrs)
if not is_eval_tuple:
eval_tsrs_to_return = eval_tsrs_to_return[0]
return eval_tsrs_to_return
else:
saved_layer[eval_tsrs[0].device] = tuple(
eval_tsr.clone() for eval_tsr in eval_tsrs)
if attribute_to_layer_input:
hook = layer.register_forward_pre_hook(forward_hook)
else:
hook = layer.register_forward_hook(forward_hook)
output = _run_forward(
forward_fn,
inputs,
target=target_ind,
additional_forward_args=additional_forward_args,
)
hook.remove()
if len(saved_layer) == 0:
raise AssertionError(
"Forward hook did not obtain any outputs for given layer")
if forward_hook_with_return:
return saved_layer, output, is_eval_tuple
return saved_layer, is_eval_tuple
