def get_target_layer(
model: Module,
layer_name: Union[str, List[str]]) -> Union[Module, List[Module]]:
if isinstance(layer_name, str):
return _get_deep_layer_name(model, layer_name)
else:
return [
_get_deep_layer_name(model, single_layer_name)
for single_layer_name in layer_name
]
def hook_removal_test_assert(self) -> None:
attr_method: Attribution
expect_error = False
if layer is not None:
if mode is HookRemovalMode.invalid_module:
expect_error = True
_set_deep_layer_value(model, layer, ErrorModule())
target_layer = _get_deep_layer_name(model, layer)
internal_algorithm = cast(Type[InternalAttribution], algorithm)
attr_method = internal_algorithm(model, target_layer)
else:
attr_method = algorithm(model)
if noise_tunnel:
attr_method = NoiseTunnel(attr_method)
if mode is HookRemovalMode.incorrect_target_or_neuron:
# Overwriting target and neuron index arguments to
# incorrect values.
if "target" in args:
args["target"] = (9999, ) * 20
expect_error = True
if "neuron_index" in args:
args["neuron_index"] = (9999, ) * 20
expect_error = True
if expect_error:
with self.assertRaises(AssertionError):
attr_method.attribute(**args)
else:
attr_method.attribute(**args)
def check_leftover_hooks(module):
self.assertEqual(len(module._forward_hooks), 0)
self.assertEqual(len(module._backward_hooks), 0)
self.assertEqual(len(module._forward_pre_hooks), 0)
model.apply(check_leftover_hooks)
def data_parallel_test_assert(self) -> None:
# Construct cuda_args, moving all tensor inputs in args to CUDA device
cuda_args = {}
for key in args:
if isinstance(args[key], Tensor):
cuda_args[key] = args[key].cuda()
elif isinstance(args[key], tuple):
cuda_args[key] = tuple(
elem.cuda() if isinstance(elem, Tensor) else elem
for elem in args[key])
else:
cuda_args[key] = args[key]
alt_device_ids = None
cuda_model = copy.deepcopy(model).cuda()
# Initialize models based on DataParallelCompareMode
if mode is DataParallelCompareMode.cpu_cuda:
model_1, model_2 = model, cuda_model
args_1, args_2 = args, cuda_args
elif mode is DataParallelCompareMode.data_parallel_default:
model_1, model_2 = (
cuda_model,
torch.nn.parallel.DataParallel(cuda_model),
)
args_1, args_2 = cuda_args, cuda_args
elif mode is DataParallelCompareMode.data_parallel_alt_dev_ids:
alt_device_ids = [0] + [
x for x in range(torch.cuda.device_count() - 1, 0, -1)
]
model_1, model_2 = (
cuda_model,
torch.nn.parallel.DataParallel(cuda_model,
device_ids=alt_device_ids),
)
args_1, args_2 = cuda_args, cuda_args
else:
raise AssertionError(
"DataParallel compare mode type is not valid.")
attr_method_1: Attribution
attr_method_2: Attribution
if target_layer:
internal_algorithm = cast(Type[InternalAttribution], algorithm)
attr_method_1 = internal_algorithm(
model_1, _get_deep_layer_name(model_1, target_layer))
# cuda_model is used to obtain target_layer since DataParallel
# adds additional wrapper.
# model_2 is always either the CUDA model itself or DataParallel
if alt_device_ids is None:
attr_method_2 = internal_algorithm(
model_2, _get_deep_layer_name(cuda_model,
target_layer))
else:
# LayerDeepLift and LayerDeepLiftShap do not take device ids
# as a parameter, since they must always have the DataParallel
# model object directly.
# Some neuron methods and GuidedGradCAM also require the
# model and cannot take a forward function.
if issubclass(
internal_algorithm,
(
LayerDeepLift,
LayerDeepLiftShap,
NeuronDeepLift,
NeuronDeepLiftShap,
NeuronDeconvolution,
NeuronGuidedBackprop,
GuidedGradCam,
),
):
attr_method_2 = internal_algorithm(
model_2,
_get_deep_layer_name(cuda_model, target_layer))
else:
attr_method_2 = internal_algorithm(
model_2.forward,
_get_deep_layer_name(cuda_model, target_layer),
device_ids=alt_device_ids,
)
else:
attr_method_1 = algorithm(model_1)
attr_method_2 = algorithm(model_2)
if noise_tunnel:
attr_method_1 = NoiseTunnel(attr_method_1)
attr_method_2 = NoiseTunnel(attr_method_2)
if attr_method_1.has_convergence_delta():
attributions_1, delta_1 = attr_method_1.attribute(
return_convergence_delta=True, **args_1)
self.setUp()
attributions_2, delta_2 = attr_method_2.attribute(
return_convergence_delta=True, **args_2)
assertTensorTuplesAlmostEqual(self,
attributions_1,
attributions_2,
mode="max",
delta=dp_delta)
assertTensorTuplesAlmostEqual(self,
delta_1,
delta_2,
mode="max",
delta=dp_delta)
else:
attributions_1 = attr_method_1.attribute(**args_1)
self.setUp()
attributions_2 = attr_method_2.attribute(**args_2)
assertTensorTuplesAlmostEqual(self,
attributions_1,
attributions_2,
mode="max",
delta=dp_delta)
def make_single_target_test(
cls,
algorithm: Type[Attribution],
model: Module,
layer: Optional[str],
args: Dict[str, Any],
target_delta: float,
noise_tunnel: bool,
baseline_distr: bool,
) -> Callable:
"""
This method creates a single target test for the given algorithm and parameters.
"""
target_layer = _get_deep_layer_name(
model, layer) if layer is not None else None
# Obtains initial arguments to replace with each example
# individually.
original_inputs = args["inputs"]
original_targets = args["target"]
original_additional_forward_args = (_format_additional_forward_args(
args["additional_forward_args"]) if "additional_forward_args"
in args else None)
num_examples = (len(original_inputs) if isinstance(
original_inputs, Tensor) else len(original_inputs[0]))
replace_baselines = "baselines" in args and not baseline_distr
if replace_baselines:
original_baselines = args["baselines"]
def target_test_assert(self) -> None:
attr_method: Attribution
if target_layer:
internal_algorithm = cast(Type[InternalAttribution], algorithm)
attr_method = internal_algorithm(model, target_layer)
else:
attr_method = algorithm(model)
if noise_tunnel:
attr_method = NoiseTunnel(attr_method)
attributions_orig = attr_method.attribute(**args)
for i in range(num_examples):
args["target"] = (original_targets[i] if len(original_targets)
== num_examples else original_targets)
args["inputs"] = (original_inputs[i:i + 1] if isinstance(
original_inputs, Tensor) else tuple(
original_inp[i:i + 1]
for original_inp in original_inputs))
if original_additional_forward_args is not None:
args["additional_forward_args"] = tuple(
single_add_arg[i:i + 1] if isinstance(
single_add_arg, Tensor) else single_add_arg
for single_add_arg in original_additional_forward_args)
if replace_baselines:
if isinstance(original_inputs, Tensor):
args["baselines"] = original_baselines[i:i + 1]
elif isinstance(original_baselines, tuple):
args["baselines"] = tuple(
single_baseline[i:i + 1] if isinstance(
single_baseline, Tensor) else single_baseline
for single_baseline in original_baselines)
self.setUp()
single_attr = attr_method.attribute(**args)
current_orig_attributions = (
attributions_orig[i:i + 1] if isinstance(
attributions_orig, Tensor) else tuple(
single_attrib[i:i + 1]
for single_attrib in attributions_orig))
assertTensorTuplesAlmostEqual(
self,
current_orig_attributions,
single_attr,
delta=target_delta,
mode="max",
)
if len(original_targets) == num_examples:
# If original_targets contained multiple elements, then
# we also compare with setting targets to a list with
# a single element.
args["target"] = original_targets[i:i + 1]
self.setUp()
single_attr_target_list = attr_method.attribute(**args)
assertTensorTuplesAlmostEqual(
self,
current_orig_attributions,
single_attr_target_list,
delta=target_delta,
mode="max",
)
return target_test_assert
