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")