Batches in guided_backprop.py

[MoH-assan]

Does guided_backprop.py support batchs?

I think the current implementation is hard-coded of a single (2D) image [1,C,W,H].

pytorch-grad-cam/pytorch_grad_cam/guided_backprop.py

Line 87 in 18144f2

loss = output[0, target_category]

I think this change is needed to support batches

    def __call__(self, input_img, target_category=None):
        replace_all_layer_type_recursive(self.model,
                                         torch.nn.ReLU,
                                         GuidedBackpropReLUasModule())


        input_img = input_img.to(self.device)

        input_img = input_img.requires_grad_(True)

        output = self.forward(input_img)

        target_category = output.argmax(dim=-1).detach() if target_category is None else target_category


        #loss = output[0, target_category]
        #loss.backward(retain_graph=True)
        loss = torch.sum(output[torch.arange(output.size(0)), target_category])
        output = torch.autograd.grad(loss, input_img, create_graph=True)[0] # Batch x Channel x Height x Width
        
        # output = input_img.grad.cpu().data.numpy()
        # output = output[0, :, :, :]
        # output = output.transpose((1, 2, 0))

        replace_all_layer_type_recursive(self.model,
                                         GuidedBackpropReLUasModule,
                                         torch.nn.ReLU())

        return output

Thanks,
Mohamed