Gradient Activation Maps for Regression and the .backward() function

Hi there,

I have a theoretical question about the .backward() function when it is computed on the output tensor rather than the loss in the context of creating activation maps from a regression problem.

For some context, I think that I understand what happens when we have class activation maps in the following scenario. For this model, I output a vector with logits for each class for each sample (before sending through to softmax) and in order to backpropagate gradients only wrt the label class, there is a one-hot encoded vector that essentially only backpropagates the signal from the output encoded for the label:

class GradCAM(object):
    """Preformatted text
    "Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization"
    https://arxiv.org/pdf/1610.02391.pdf
    Look at Figure 2 on page 4
    """

    def __init__(self, model, candidate_layers=None):
        super(GradCAM, self).__init__()
        self.model = model
        self.handlers = []  # a set of hook function handlers
        self.fmap_pool = OrderedDict()
        self.grad_pool = OrderedDict()
        self.candidate_layers = candidate_layers  # list

        def forward_hook(module, input, output):
            # Save featuremaps
            self.fmap_pool[id(module)] = output.detach()

        def backward_hook(module, grad_in, grad_out):
            # Save the gradients correspond to the featuremaps
            self.grad_pool[id(module)] = grad_out[0].detach()

        # If any candidates are not specified, the hook is registered to all the layers.
        for module in self.model.named_modules():
            if self.candidate_layers is None or module[0] in self.candidate_layers:
                self.handlers.append(module[1].register_forward_hook(forward_hook))
                self.handlers.append(module[1].register_backward_hook(backward_hook))
                
    def _encode_one_hot(self, ids):
        one_hot = torch.zeros_like(self.logits).cuda()
        one_hot.scatter_(1, ids, 1.0)
        return one_hot

    def _find(self, pool, target_layer):
        for key, value in pool.items():
            for module in self.model.named_modules():
                if id(module[1]) == key:
                    if module[0] == target_layer:
                        return value
        raise ValueError("Invalid layer name: {}".format(target_layer))

    def _compute_grad_weights(self, grads):
        return F.adaptive_avg_pool2d(grads, 1)
    
    def forward(self, image):
        """
        multi-label classification
        """
        self.image_shape = image.shape[2:]
        self.model.zero_grad()
        self.logits = self.model(image)
        self.probs = F.sigmoid(self.logits)
        return self.probs.sort(dim=1, descending=True)
    
    def backward(self, one_hot):
        """
        Class-specific backpropagation
        Either way works:
        1. self.logits.backward(gradient=one_hot, retain_graph=True)
        2. (self.logits * one_hot).sum().backward(retain_graph=True)
        """

        #one_hot = self._encode_one_hot(ids)
        self.logits.backward(gradient=one_hot, retain_graph=True)

    def generate(self, target_layer):
        fmaps = self._find(self.fmap_pool, target_layer)
        grads = self._find(self.grad_pool, target_layer)
        weights = self._compute_grad_weights(grads)

        gcam = torch.mul(fmaps, weights).sum(dim=1, keepdim=True)
        gcam = F.relu(gcam)

        gcam = F.interpolate(
            gcam, self.image_shape, mode="bilinear", align_corners=False
        )

        B, C, H, W = gcam.shape
        gcam = gcam.view(B, -1)
        gcam -= gcam.min(dim=1, keepdim=True)[0]
        gcam /= gcam.max(dim=1, keepdim=True)[0]
        gcam = gcam.view(B, C, H, W)
        return gcam
    
    def remove_hook(self):
        """
        Remove all the forward/backward hook functions
        """
        for handle in self.handlers:
            handle.remove()

However, I get a bit more confused when it comes to regression. In my regression models, I output a single number after a fully connected layer; therefore, I only have one number, not a vector of logits. Is it then reasonable to use a tensor of ones in place of a one-hot encoded vector in order to backpropagate wrt to the singular logit output? Essentially, putting in a vector of 1s inside the backpropagation algorithm E.g.:

class GradAM(object):
    """
    "GradAM: Adapted from GradCAM, does not need class-specific activations for regression or binary classification" 
    """

    def __init__(self, model, loader, candidate_layers=None):
        super(GradAM, self).__init__()
        self.model = model
        self.handlers = []  # a set of hook function handlers
        self.fmap_pool = OrderedDict()
        self.grad_pool = OrderedDict()
        self.candidate_layers = candidate_layers  # list
        
        def forward_hook(module, input, output):
            # Save featuremaps
            self.fmap_pool[id(module)] = output.detach()

        def backward_hook(module, grad_in, grad_out):
            # Save the gradients correspond to the featuremaps
            self.grad_pool[id(module)] = grad_out[0].detach()

        # If any candidates are not specified, the hook is registered to all the layers.
        for module in self.model.named_modules():
            if self.candidate_layers is None or module[0] in self.candidate_layers:
                self.handlers.append(module[1].register_forward_hook(forward_hook))
                self.handlers.append(module[1].register_backward_hook(backward_hook))
                

    def _find(self, pool, target_layer):
        for key, value in pool.items():
            for module in self.model.named_modules():
                if id(module[1]) == key:
                    if module[0] == target_layer:
                        return value
        raise ValueError("Invalid layer name: {}".format(target_layer))

    def _compute_grad_weights(self, grads):
        return F.adaptive_avg_pool2d(grads, 1)
    
    def forward(self, image):
        self.image_shape = image.shape[2:]
        self.model.zero_grad()
        self.outputs = self.model(image)
        return self.outputs

        
    def backward(self): 
        self.outputs.backward(torch.ones_like(self.outputs), retain_graph = True)

    def generate(self, target_layer):
        fmaps = self._find(self.fmap_pool, target_layer)
        grads = self._find(self.grad_pool, target_layer)
        weights = self._compute_grad_weights(grads)

        gcam = torch.mul(fmaps, weights).sum(dim=1, keepdim=True)
        gcam = F.relu(gcam)

        gcam = F.interpolate(
            gcam, self.image_shape, mode="bilinear", align_corners=False
        )

        B, C, H, W = gcam.shape
        gcam = gcam.view(B, -1)
        gcam -= gcam.min(dim=1, keepdim=True)[0]
        gcam /= gcam.max(dim=1, keepdim=True)[0]
        gcam = gcam.view(B, C, H, W)
        return gcam
    
    def remove_hook(self):
        """
        Remove all the forward/backward hook functions
        """
        for handle in self.handlers:
            handle.remove()

(My apologies that some of the variables have been changed but the basic structure of the classes is still quite comparable).

Or instead, does it make more sense to backpropagate wrt to the loss – e.g. just calculating the loss from these outputs and then using the .backward() function default parameters with the exception of adding retains_graph = True?

In this publication regarding regression activation mapping, they do something similar to what I’m talking about by doing a riff on CAM. However, this implies that you must still have a GAP layer to determine the weights (e.g. they weights of the GAP layer wrt to the final output are now the weights for the weighted average of the final convolutional layer as well), but I would like to extend the concept to any NN structure and not rely on having a GAP layer.

Thanks for reading and any discussion is appreciated.