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 jit_test_assert(self) -> None:
model_1 = model
attr_args = args
if (mode is JITCompareMode.data_parallel_jit_trace
or JITCompareMode.data_parallel_jit_script):
if not torch.cuda.is_available() or torch.cuda.device_count(
) == 0:
raise unittest.SkipTest(
"Skipping GPU test since CUDA not available.")
# 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]
attr_args = cuda_args
model_1 = model_1.cuda()
# Initialize models based on JITCompareMode
if (mode is JITCompareMode.cpu_jit_script
or JITCompareMode.data_parallel_jit_script):
model_2 = torch.jit.script(model_1) # type: ignore
elif (mode is JITCompareMode.cpu_jit_trace
or JITCompareMode.data_parallel_jit_trace):
all_inps = _format_input(args["inputs"]) + (
_format_additional_forward_args(
args["additional_forward_args"])
if "additional_forward_args" in args and
args["additional_forward_args"] is not None else tuple())
model_2 = torch.jit.trace(model_1, all_inps) # type: ignore
else:
raise AssertionError("JIT compare mode type is not valid.")
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, **attr_args)
self.setUp()
attributions_2, delta_2 = attr_method_2.attribute(
return_convergence_delta=True, **attr_args)
assertTensorTuplesAlmostEqual(self,
attributions_1,
attributions_2,
mode="max")
assertTensorTuplesAlmostEqual(self,
delta_1,
delta_2,
mode="max")
else:
attributions_1 = attr_method_1.attribute(**attr_args)
self.setUp()
attributions_2 = attr_method_2.attribute(**attr_args)
assertTensorTuplesAlmostEqual(self,
attributions_1,
attributions_2,
mode="max")