Cam for multi gpu training

[BELZHANG]

Hi, I'm trying to use heatmap during training. it's ok when I process the cam using one GPU, but when I do the multi-GPU training, the default cam setting doesn't seem to work. RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument weight in method wrapper_cudnn_convolution) I try to data parallel the cam like this cam_target_layer = [model.module.Res50.layer4[-1]] grad_cam = GradCAM(model=nn.DataParallel(model.module.Res50), target_layers=cam_target_layer, use_cuda=True) then the output batch size becomes half the input, like input(16, 3, 224, 224) the output cam (8, 7, 7) Is it possible for the grad_cam process parallel on multi-GPU, if so, how to set up the cam correctly?

Thank you!

[DarkLeaves]

the same question

[xinyanghuang7]

Hello, have you resolved the issue?

[DarkLeaves]

not yet

[memory009]

base_camhas a model.cuda() operation, so you new to add a new parameter likes device when the class BASE_CAM was init.

This is how I modified it:

class BaseCAM:
    def __init__(self,
                 model: torch.nn.Module,
                 target_layers: List[torch.nn.Module],
                 use_cuda: bool = False,
                 reshape_transform: Callable = None,
                 compute_input_gradient: bool = False,
                 uses_gradients: bool = True,
                 # add device as new parameter
                 device: torch.device = None) -> None:
        self.model = model.eval()
        self.target_layers = target_layers
        self.cuda = use_cuda
        # Assign the value of the `device` parameter to the member variable `self.device` of the class
        self.device = device
        if self.cuda:
            self.model = model.cuda()
        self.reshape_transform = reshape_transform
        self.compute_input_gradient = compute_input_gradient
        self.uses_gradients = uses_gradients
        self.activations_and_grads = ActivationsAndGradients(
            self.model, target_layers, reshape_transform)

Example usage

device = torch.device("cuda:0") # Change to your desired CUDA device cam_instance = BaseCAM(model, target_layers, use_cuda=True, device=device)