Backward hook registered using register_hook function not being called during model's backward pass

netgvarun2012: <p>I am trying to implement GradCam like functionality for Large Vision Language Models (LVLMs).</p>
<p>I have implemented following code to capture gradients during backward pass and output activations during forward pass by attaching hooks to selected layer.</p>
<p>Since I am working with Large Vision language models, I have a VIT (Vision transformer) as well as a LLM decoder which is finally used to produce text output. I am trying to fetch gradient of output of attention layer w.r.t highest logit of next token predicted by the model.</p>
<p>However, because of an existing bug (<a href="https://github.com/pytorch/pytorch/issues/61519" class="inline-onebox" rel="noopener nofollow ugc">Inplace ReLU incompatible with backward hook · Issue #61519 · pytorch/pytorch · GitHub</a>), I don’t use backward hook to record gradients.</p>
<p>Gradients are still captured during the backward pass, just not through the backward hooks. Instead, during the forward pass, the code uses a <strong>register_hook</strong> function to register a function called ‘<strong>_store_grad</strong>’ that will be called during backward pass.</p>
<p>So, while the backward hook isn’t used directly due to the bug, the gradients are still captured by attaching a hook to the output tensor during the forward pass.</p>
<p>But the same is not happening, even upon explicitly calling <strong>target_logit.backward()</strong>, gradient is not getting captured as <strong>'_store_grad</strong>’ function is not getting triggered for some reason.</p>
<p>What could be the reason for function not getting triggered during backward pass?</p>
<pre><code class="lang-auto">import openai
import cv2  # OpenCV for image processing
import numpy as np
import json
import os
import torch
import json
from transformers import BitsAndBytesConfig
from transformers import AutoProcessor,LlavaForConditionalGeneration, AutoModelForCausalLM,AutoModelForPreTraining
from PIL import Image
import matplotlib.pyplot as plt

class LVLMGradCam:
    def __init__(self, model, target_layer,processor,token_num):
        self.model = model.eval()  # Set the model to evaluation mode
        self.feature = None
        self.gradient = None
        self.handlers = []
        self.target_layer = target_layer
        self.processor = processor
        self.token_num = token_num
        self._get_hook()

    def _get_features_hook(self, module, input, output):
        print("Forward hook executed to capture the activations\n")  
        self.feature = self.reshape_transform(output[0])        
    
    def _get_grads_hook(self, module, input_grad, output_grad):
        self.gradient = self.reshape_transform(output_grad[0])  # Store and reshape the output gradients

        def _store_grad(grad):
            print(f"Backward hook executed")  # Debugging line            
            self.gradient = self.reshape_transform(grad)  # Store gradients for later use
            print(f'self.gradient being stored is {self.gradient}\n')            
        
        print("Forward hook executed to register a function that will be called during the backward pass\n")  
        output_grad[0].register_hook(_store_grad)  # Register hook to store gradients

    # Register forward hooks to the target layer
    def _get_hook(self):
        print("Above registering Forward hooks\n")  # Debugging line            
        self.target_layer.register_forward_hook(self._get_features_hook)
        self.target_layer.register_forward_hook(self._get_grads_hook)        
        print('Hooks registered\n')
    
    def reshape_transform(self, inp_tensor, height=24, width=24):   
        if isinstance(inp_tensor, tuple):
            inp_tensor = inp_tensor[0]
        result = inp_tensor[:, 1:, :].reshape(inp_tensor.size(0), height, width, inp_tensor.size(2))
        
        result = result.transpose(2, 3).transpose(1, 2)  # Rearrange dimensions to (C, H, W)
        return result
    
    def __call__(self, inputs):
        print(f'self.gradient is {self.gradient}\n')
        self.model.zero_grad()
        output = self.model.generate(**inputs, max_new_tokens=50, output_scores=True, return_dict_in_generate=True,use_cache=True)
        
        # Extract the generated token IDs
        generated_token_ids = output['sequences']
        tokenizer = self.processor.tokenizer
        gen_text = processor.tokenizer.decode(generated_token_ids[0],skip_special_tokens=True)
        print(f'generated text is {gen_text}\n')
        
        # Get the input IDs for reference
        input_ids = inputs['input_ids'][0]
        
        # Access the logits for the first generated token directly
        generated_logits = output['scores'][self.token_num][0]  # All logits for the first token generation
        
        # Find the id of the first generated token
        generated_id = generated_token_ids[0][len(input_ids):][self.token_num]
        gen_text = processor.tokenizer.decode(generated_id,skip_special_tokens=True)
        print(f'generated next token is {gen_text}\n')
        
        # Get the specific logit for that token based on the id
        target_logit = generated_logits[generated_id]  
        target_logit.backward()  # Backward pass to compute gradients

        # Get gradients and features
        gradients = self.gradient.cpu().data.numpy()
        weight = np.mean(gradients, axis=(1, 2))  # Average the gradients
        feature = self.feature[0].cpu().data.numpy()

        # Compute the weighted sum of the features
        cam = feature * weight[:, np.newaxis, np.newaxis]
        cam = np.sum(cam, axis=0)  # Sum over the channels
        cam = np.maximum(cam, 0)  # Apply ReLU to remove negative values        

        # Normalize the heatmap
        cam -= np.min(cam)
        cam /= np.max(cam)

        return cam  # Return the Grad-CAM heatmap

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)


model_id = "llava-hf/llava-1.5-7b-hf"

processor = AutoProcessor.from_pretrained(model_id)
model = LlavaForConditionalGeneration.from_pretrained(model_id, quantization_config=quantization_config, device_map="auto")

input_file = '/home/HADES/dataset/instructions/Financial.json'
with open(input_file, 'r') as f:
    dataset = json.load(f)

keyword = dataset[25]['keywords']
imagefile = '/home/HADES/dataset/instructions/google_images/Financial/'+dataset[25]['image']
instruction = dataset[25]['instruction']

image1 = Image.open(imagefile)
img = cv2.imread(imagefile)  # Load the image using OpenCV


prompts = [
            "USER: &lt;image&gt;\nShow steps to print a bank statement.\nASSISTANT:"
         ]

inputs = processor(prompts, image1, padding=True, return_tensors="pt").to("cuda")
for k,v in inputs.items():
  print(k,v.shape)

target_layer =  model.vision_tower.vision_model.encoder.layers[-1].self_attn 
grad_cam = LVLMGradCam(model, target_layer,processor,9)
heatmap = grad_cam(inputs)
# Visualize the heatmap
plt.imshow(heatmap[0], cmap='jet', alpha=0.5)  
plt.axis('off')
plt.show()
</code></pre>