modified visual encoder

app.py

@@ -25,6 +25,7 @@ model_utils, vl_gpt, tokenizer = None, None, None
 model_name = "Clip"
 language_model_max_layer = 24
 language_model_best_layer = 8
+vision_model_best_layer = 24
 
 def clean():
     global model_utils, vl_gpt, tokenizer, clip_utils
@@ -116,7 +117,10 @@ def multimodal_understanding(model_type,
         if activation_map_method == "GradCAM":
             # target_layers = vl_gpt.vision_model.vision_tower.blocks
             if focus == "Visual Encoder":
-                all_layers = [block.norm1 for block in vl_gpt.vision_model.vision_tower.blocks]
+                if model_name.split('-')[0] == "Janus":
+                    all_layers = [block.norm1 for block in vl_gpt.vision_model.vision_tower.blocks]
+                else:
+                    all_layers = [block.layer_norm1 for block in vl_gpt.vision_tower.vision_model.encoder.layers]
             else:
                 all_layers = [layer.self_attn for layer in vl_gpt.language_model.model.layers]
             
@@ -137,17 +141,33 @@ def multimodal_understanding(model_type,
                 gradcam = AttentionGuidedCAMChartGemma(vl_gpt, target_layers)
 
             start = 0
+            cam = []
             if focus == "Visual Encoder":
-                cam_tensors, grid_size = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+                if target_token_idx != -1:
+                    cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+                    cam_grid = cam_tensors.reshape(grid_size, grid_size)
+                    cam_i = generate_gradcam(cam_grid, image)
+                    cam_i = add_title_to_image(cam_i, input_ids_decoded[start + target_token_idx])
+                    cam = [cam_i]
+                else:
+                    i = 0
+                    cam = []
+                    while start + i < len(input_ids_decoded):
+                        if model_name.split('-')[0] == "Janus":
+                            gradcam = AttentionGuidedCAMJanus(vl_gpt, target_layers)
+                        elif model_name.split('-')[0] == "LLaVA":
+                            gradcam = AttentionGuidedCAMLLaVA(vl_gpt, target_layers)
+                        elif model_name.split('-')[0] == "ChartGemma":
+                            gradcam = AttentionGuidedCAMChartGemma(vl_gpt, target_layers)
+                        cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_pooling_method, focus)
+                        cam_grid = cam_tensors.reshape(grid_size, grid_size)
+                        cam_i = generate_gradcam(cam_grid, image)
+                        cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
+                        cam.append(cam_i)
+                        gradcam.remove_hooks()
+                        i += 1
             else:
                 cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
-            gradcam.remove_hooks()
-
-
-            if focus == "Visual Encoder":
-                cam_grid = cam_tensors.reshape(grid_size, grid_size)
-                cam = [generate_gradcam(cam_grid, image)]
-            else:
                 if target_token_idx != -1:
                     input_text_decoded = input_ids_decoded[start + target_token_idx]
                     for i, cam_tensor in enumerate(cam_tensors):
@@ -164,6 +184,9 @@ def multimodal_understanding(model_type,
                         cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
 
                         cam.append(cam_i)
+                     
+            gradcam.remove_hooks()
+                
 
     # Collect Results
     RESULTS_ROOT = "./results"
@@ -193,7 +216,7 @@ def multimodal_understanding(model_type,
 # Gradio interface
 
 def model_slider_change(model_type):
-    global model_utils, vl_gpt, tokenizer, clip_utils, model_name, language_model_max_layer, language_model_best_layer
+    global model_utils, vl_gpt, tokenizer, clip_utils, model_name, language_model_max_layer, language_model_best_layer, vision_model_best_layer
     model_name = model_type
     if model_type == "Clip":
         clean()
@@ -251,13 +274,14 @@ def model_slider_change(model_type):
         model_utils = ChartGemma_Utils()
         vl_gpt, tokenizer = model_utils.init_ChartGemma()
         language_model_max_layer = 18
+        vision_model_best_layer = 19
         language_model_best_layer = 15
 
         res = (
             gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
             gr.Slider(minimum=1, maximum=language_model_best_layer, value=language_model_best_layer, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=language_model_best_layer, value=language_model_best_layer, step=1, label="visualization layers max"),
-            gr.Dropdown(choices=["Language Model"], value="Language Model", label="focus"),
+            gr.Dropdown(choices=["Visual Encoder", "Language Model"], value="Language Model", label="focus"),
             gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
         )
         return res
@@ -292,12 +316,21 @@ def focus_change(focus):
             return res
 
     else:
-        res = (
-            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-            gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers min"), 
-            gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers max")
-        )
-        return res
+        if model_name.split('-')[0] == "ChartGemma":
+            res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers min"), 
+                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers max")
+            )
+            return res
+        
+        else:
+            res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers min"), 
+                gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers max")
+            )
+            return res
 
 
 
@@ -305,27 +338,37 @@ def focus_change(focus):
 
 with gr.Blocks() as demo:
     gr.Markdown(value="# Multimodal Understanding")
+    
+    with gr.Row():
+        image_input = gr.Image(height=500, label="Image")
+        activation_map_output = gr.Gallery(label="Visualization", height=500, columns=1, preview=True)
+
+    with gr.Row():
+        chart_type = gr.Textbox(label="Chart Type")
+        understanding_output = gr.Textbox(label="Answer")
+
     with gr.Row():
-        with gr.Column():
-            image_input = gr.Image()
-            activation_map_output = gr.Gallery(label="activation Map", height=300, columns=1)
 
         with gr.Column():
             model_selector = gr.Dropdown(choices=["Clip", "ChartGemma-3B", "Janus-Pro-1B", "Janus-Pro-7B", "LLaVA-1.5-7B"], value="Clip", label="model")
+            question_input = gr.Textbox(label="Input Prompt")
+            und_seed_input = gr.Number(label="Seed", precision=0, value=42)
+            top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
+            temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
+            target_token_idx = gr.Number(label="target_token_idx (-1 means all)", precision=0, value=-1)
+
+        
+        with gr.Column():
             response_type = gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type")
             focus = gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus")
-            activation_map_method = gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
+            activation_map_method = gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="visualization type")
             visual_pooling_method = gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
             
 
             visualization_layers_min = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min")
             visualization_layers_max = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max")
+
         
-            question_input = gr.Textbox(label="Question")
-            und_seed_input = gr.Number(label="Seed", precision=0, value=42)
-            top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
-            temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
-            target_token_idx = gr.Number(label="target_token_idx (-1 means all)", precision=0, value=-1)
         
 
 
@@ -360,8 +403,7 @@ with gr.Blocks() as demo:
         
 
     understanding_button = gr.Button("Submit")
-    chart_type = gr.Textbox(label="Chart Type")
-    understanding_output = gr.Textbox(label="Answer")
+    
     understanding_target_token_decoded_output = gr.Textbox(label="Target Token Decoded")
 
 

demo/cam.py

@@ -85,10 +85,11 @@ class AttentionGuidedCAMClip(AttentionGuidedCAM):
             print("act shape", act.shape)
             print("grad_weights shape", grad_weights.shape)
             
-            # cam = (act * grad_weights).sum(dim=-1)  # Weighted activation map
+            # cam = (act * grad_weights).sum(dim=-1)
             cam, _ = (act * grad_weights).max(dim=-1)
+            # cam, _ = act.max(dim=-1)
+            # cam = cam.unsqueeze(0)
             # cam, _ = grad_weights.max(dim=-1)
-            # cam = self.normalize(cam)
             print("cam_shape: ", cam.shape)
 
             # Sum across all layers
@@ -166,20 +167,23 @@ class AttentionGuidedCAMJanus(AttentionGuidedCAM):
 
         if focus == "Visual Encoder":
             # Pooling
-            if visual_pooling_method == "CLS":
-                image_embeddings_pooled = image_embeddings[:, 0, :]
-            elif visual_pooling_method == "avg":
-                image_embeddings_pooled = image_embeddings[:, 1:, :].mean(dim=1) # end of image: 618
-            elif visual_pooling_method == "max":
-                image_embeddings_pooled, _ = image_embeddings[:, 1:, :].max(dim=1)
-
-            print("image_embeddings_shape: ", image_embeddings_pooled.shape)
-            
+            # if visual_pooling_method == "CLS":
+            #     image_embeddings_pooled = image_embeddings[:, 0, :]
+            # elif visual_pooling_method == "avg":
+            #     image_embeddings_pooled = image_embeddings[:, 1:, :].mean(dim=1)
+            # elif visual_pooling_method == "max":
+            #     image_embeddings_pooled, _ = image_embeddings[:, 1:, :].max(dim=1)
 
+            # print("image_embeddings_shape: ", image_embeddings_pooled.shape)
+            
 
-            inputs_embeddings_pooled = inputs_embeddings[:, 620: -4].mean(dim=1)
+            start_idx = 620
+            # inputs_embeddings_pooled = inputs_embeddings[:, start_idx: -4].mean(dim=1)
             self.model.zero_grad()
-            image_embeddings_pooled.backward(inputs_embeddings_pooled, retain_graph=True)
+            # image_embeddings_pooled.backward(inputs_embeddings_pooled, retain_graph=True)
+
+            loss = outputs.logits.max(dim=-1).values[0, start_idx + class_idx]
+            loss.backward()
 
             cam_sum = None
             for act, grad in zip(self.activations, self.gradients):
@@ -195,6 +199,7 @@ class AttentionGuidedCAMJanus(AttentionGuidedCAM):
                 print("grad_weights shape", grad_weights.shape)
 
                 cam, _ = (act * grad_weights).max(dim=-1)
+                # cam, _ = grad_weights.max(dim=-1)
                 print(cam.shape)
 
                 # Sum across all layers
@@ -224,7 +229,7 @@ class AttentionGuidedCAMJanus(AttentionGuidedCAM):
             cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
             cam_sum = cam_sum.detach().to("cpu")
 
-            return cam_sum, grid_size
+            return cam_sum, grid_size, start_idx
 
 
 
@@ -407,7 +412,7 @@ class AttentionGuidedCAMLLaVA(AttentionGuidedCAM):
 class AttentionGuidedCAMChartGemma(AttentionGuidedCAM):
     def __init__(self, model, target_layers):
         self.target_layers = target_layers
-        super().__init__(model, register=False)
+        super().__init__(model, register=True)
         self._modify_layers()
         self._register_hooks_activations()
     
@@ -445,12 +450,9 @@ class AttentionGuidedCAMChartGemma(AttentionGuidedCAM):
             for param in layer.parameters():
                 param.requires_grad = True
         
-        outputs_raw = self.model(**inputs)
+        outputs_raw = self.model(**inputs, output_hidden_states=True)
+
 
-        self.model.zero_grad()
-        # print(outputs_raw)
-        loss = outputs_raw.logits.max(dim=-1).values.sum()
-        loss.backward()
 
         # get image masks
         image_mask = []
@@ -462,61 +464,135 @@ class AttentionGuidedCAMChartGemma(AttentionGuidedCAM):
                 last = i
             else:
                 image_mask.append(False)
+        start_idx = last + 1
 
 
-        # Aggregate activations and gradients from ALL layers
-        self.activations = [layer.get_attn_map() for layer in self.target_layers]
-        self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
-        print(f"layers shape: {len(self.target_layers)}")
-        print("activations & gradients shape", len(self.activations), len(self.gradients))
 
-        cams = []
-        
-        # Ver 2
-        for act, grad in zip(self.activations, self.gradients):
-
-            print("act shape", act.shape)
-            print("grad shape", grad.shape)
+        if focus == "Visual Encoder":
+            # image_embeddings = outputs_raw.image_hidden_states
+            # inputs_embeddings = outputs_raw.hidden_states[0]
+            # # Pooling
+            # if visual_pooling_method == "avg":
+            #     image_embeddings_pooled = image_embeddings.mean(dim=1) # end of image: 618
+            # elif visual_pooling_method == "max":
+            #     image_embeddings_pooled, _ = image_embeddings.max(dim=1)
+
+            # print("image_embeddings_shape: ", image_embeddings_pooled.shape)
             
-            grad = F.relu(grad)
 
-            cam = act * grad # shape: [1, heads, seq_len, seq_len]
-            cam = cam.sum(dim=1) # shape: [1, seq_len, seq_len]
-            cam = cam.to(torch.float32).detach().cpu()
-            cams.append(cam)
 
-        # cam_sum = F.relu(cam_sum)
-        # cam_sum = cam_sum.to(torch.float32)
+            # inputs_embeddings_pooled = inputs_embeddings[:, start_idx:].mean(dim=1)
+            self.model.zero_grad()
+            # image_embeddings_pooled.backward(inputs_embeddings_pooled, retain_graph=True)
 
-        # cams shape: [layers, 1, seq_len, seq_len]
-        cam_sum_lst = []
+            loss = outputs_raw.logits.max(dim=-1).values[0, start_idx + class_idx]
+            loss.backward()
 
-        start_idx = last + 1
-        for i in range(start_idx, cams[0].shape[1]):
             cam_sum = None
-            for layer, cam_l in enumerate(cams):
-                cam_l_i = cam_l[0, i, :] # shape: [1: seq_len]
+            for act, grad in zip(self.activations, self.gradients):
+                # act = torch.sigmoid(act)
+                act = F.relu(act[0])
+    
 
-                cam_l_i = cam_l_i[image_mask].unsqueeze(0) # shape: [1, img_seq_len]
-                # print(f"layer: {layer}, token index: {i}")
-                # print("cam_sum shape: ", cam_l_i.shape)
-                num_patches = cam_l_i.shape[-1]  # Last dimension of CAM output
-                grid_size = int(num_patches ** 0.5)
-                # print(f"Detected grid size: {grid_size}x{grid_size}")
+                # Compute mean of gradients
+                print("grad shape:", grad.shape)
+                grad_weights = grad.mean(dim=-1, keepdim=True)
 
-                # Fix the reshaping step dynamically
-                cam_reshaped = cam_l_i.view(grid_size, grid_size)
-                # print(f"max: {cam_reshaped.max()}, min: {cam_reshaped.min()}")
-                cam_normalized = (cam_reshaped - cam_reshaped.min()) / (cam_reshaped.max() - cam_reshaped.min())
-                if cam_sum == None:
-                    cam_sum = cam_normalized
+                print("act shape", act.shape)
+                print("grad_weights shape", grad_weights.shape)
+
+                cam = (act * grad_weights).sum(dim=-1)
+                # cam, _ = (act * grad_weights).max(dim=-1)
+                # cam, _ = grad_weights.max(dim=-1)
+                print(cam.shape)
+
+                # Sum across all layers
+                if cam_sum is None:
+                    cam_sum = cam
                 else:
-                    cam_sum += cam_normalized
-                # print(f"normalized: max: {cam_normalized.max()}, min: {cam_normalized.min()}")
+                    cam_sum += cam  
+
+            # Normalize
+            cam_sum = F.relu(cam_sum)
+            
+
+            # thresholding
+            cam_sum = cam_sum.to(torch.float32).detach().cpu()
+            percentile = torch.quantile(cam_sum, 0.2)  # Adjust threshold dynamically
+            cam_sum[cam_sum < percentile] = 0
 
-            # print(f"sum: max: {cam_sum.max()}, min: {cam_sum.min()}")
+            # Reshape
+            print("cam_sum shape: ", cam_sum.shape)
+            num_patches = cam_sum.shape[-1]  # Last dimension of CAM output
+            grid_size = int(num_patches ** 0.5)
+            print(f"Detected grid size: {grid_size}x{grid_size}")
+            
+            cam_sum = cam_sum.view(grid_size, grid_size)
             cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
-            cam_sum_lst.append(cam_sum)
+
+            return cam_sum, grid_size, start_idx
+
+        elif focus == "Language Model":
+            self.model.zero_grad()
+            # print(outputs_raw)
+            loss = outputs_raw.logits.max(dim=-1).values.sum()
+            loss.backward()
+
+
+
+            # Aggregate activations and gradients from ALL layers
+            self.activations = [layer.get_attn_map() for layer in self.target_layers]
+            self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
+            print(f"layers shape: {len(self.target_layers)}")
+            print("activations & gradients shape", len(self.activations), len(self.gradients))
+
+            cams = []
+            
+            # Ver 2
+            for act, grad in zip(self.activations, self.gradients):
+
+                print("act shape", act.shape)
+                print("grad shape", grad.shape)
+                
+                grad = F.relu(grad)
+
+                cam = act * grad # shape: [1, heads, seq_len, seq_len]
+                cam = cam.sum(dim=1) # shape: [1, seq_len, seq_len]
+                cam = cam.to(torch.float32).detach().cpu()
+                cams.append(cam)
+
+            # cam_sum = F.relu(cam_sum)
+            # cam_sum = cam_sum.to(torch.float32)
+
+            # cams shape: [layers, 1, seq_len, seq_len]
+            cam_sum_lst = []
+
+            start_idx = last + 1
+            for i in range(start_idx, cams[0].shape[1]):
+                cam_sum = None
+                for layer, cam_l in enumerate(cams):
+                    cam_l_i = cam_l[0, i, :] # shape: [1: seq_len]
+
+                    cam_l_i = cam_l_i[image_mask].unsqueeze(0) # shape: [1, img_seq_len]
+                    # print(f"layer: {layer}, token index: {i}")
+                    # print("cam_sum shape: ", cam_l_i.shape)
+                    num_patches = cam_l_i.shape[-1]  # Last dimension of CAM output
+                    grid_size = int(num_patches ** 0.5)
+                    # print(f"Detected grid size: {grid_size}x{grid_size}")
+
+                    # Fix the reshaping step dynamically
+                    cam_reshaped = cam_l_i.view(grid_size, grid_size)
+                    # print(f"max: {cam_reshaped.max()}, min: {cam_reshaped.min()}")
+                    # cam_reshaped = (cam_reshaped - cam_reshaped.min()) / (cam_reshaped.max() - cam_reshaped.min())
+                    if cam_sum == None:
+                        cam_sum = cam_reshaped
+                    else:
+                        cam_sum += cam_reshaped
+                    # print(f"normalized: max: {cam_normalized.max()}, min: {cam_normalized.min()}")
+
+                # print(f"sum: max: {cam_sum.max()}, min: {cam_sum.min()}")
+                cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
+                cam_sum_lst.append(cam_sum)
 
 
         return cam_sum_lst, grid_size, start_idx