Inquiry about Attention Visualization in MambaVision

[mason0417]

Your recent work on MambaVision has left a profound impact on me. The way you seamlessly integrate advanced techniques to achieve state-of-the-art results in image classification, object detection, and segmentation is nothing short of remarkable. It’s clear that your expertise and creativity are driving significant advancements in our field.
I have been particularly intrigued by the attention visualization results you presented in your recent publication. It’s fascinating to see how the model’s attention mechanisms can be so effectively visualized, providing deep insights into the model’s decision-making process. I have been trying to implement similar attention visualization heatmaps in my own research, but unfortunately, I have not yet succeeded.
Upon reviewing the open-source code, I noticed that the specific implementation details for attention visualization are not included. This has left me with some questions about how to extract the attention weights and visualize them effectively. I understand that you must be extremely busy, but I would be incredibly grateful if you could share some insights or guidance on this topic. Your expertise would be invaluable to me as I continue to explore and advance my research in this area.

[ahatamiz]

Hi @mason0417

Thank you for the kind words.

We did not (and still dont) officially support visualization of attention maps in order to keep the code clean and consistent with its previous releases.

But we simply followed very standard approaches for visualizing the attention map. If you extract the features from any of the self-attention layers (preferably the last ones) of stages 3 or 4, then you can do a binary interpolation to the image size, simply normalize the score map and overlay it on top of the original image. For normalizing score map, you can do something like the following:

import numpy as np

def normalize_scoremap(map):
    if map.min() == map.max():
        return np.zeros_like(map)
    map -= map.min()
    map /= map.max()
    return map

Hope this is useful.

[mason0417]

I'm working on visualizing attention maps from all Attention layers in the MambaVision-T-1K model, and I've written some code to achieve this. I'd like to verify if my approach is correct.
In my code, I've registered hooks to capture the output of each Attention layer during the forward pass. I then process each attention map by removing the batch dimension, averaging across channels if necessary, and reshaping it into a 2D feature map. After resizing the attention map to match the original image size using bilinear interpolation, I normalize it to ensure the values are between 0 and 1. Finally, I visualize the attention maps for each layer and create a combined attention map by summing all layer attention maps and normalizing the result.
Could you please review my approach and let me know if there are any improvements or corrections needed? I appreciate your help!
Best regards,
jiajun peng

My code：
import torch
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageClassification
from google.colab import files
import torch.nn.functional as F

Hook function to capture attention maps

def get_attention_hook(module, input, output):
print(f"Hook triggered in module: {module.class.name}")
print(f"Output shape: {output.shape}")
attention_maps.append(output)

List to store attention maps

attention_maps = []

Register hooks on attention layers

def register_attention_hooks(model):
hooks = []
for name, module in model.named_modules():
if module.class.name == "Attention":
hook = module.register_forward_hook(get_attention_hook)
hooks.append(hook)
return hooks

Normalization function for score maps

def normalize_scoremap(map):
if map.min() == map.max():
return np.zeros_like(map)
map -= map.min()
map /= map.max()
return map

Load pre-trained model

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
model = model.to(device)
model.eval()

Register hooks

hooks = register_attention_hooks(model)

Upload image

uploaded = files.upload()
img_path = list(uploaded.keys())[0]
original_img = Image.open(img_path).convert("RGB") # Save the original image
img = original_img.copy()

Preprocess the image

transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
img_tensor = transform(img).unsqueeze(0).to(device)

Forward pass to capture attention maps

with torch.no_grad():
output = model(img_tensor)

Process and visualize attention maps from all layers

if attention_maps:
# Get the original image dimensions
original_width, original_height = original_img.size

# Create a figure to display attention maps from all layers
num_layers = len(attention_maps)
fig, axs = plt.subplots(num_layers + 1, 2, figsize=(15, 5 * (num_layers + 1)))

# Store the combined attention map
combined_attn_map = np.zeros((original_height, original_width))

for i, attn_map in enumerate(attention_maps):
    # Process the current layer's attention map
    current_attn_map = attn_map.squeeze(0)  # Remove batch dimension

    # Average across channels if necessary
    if len(current_attn_map.shape) > 1:
        current_attn_map = current_attn_map.mean(dim=1)  # Shape: [sequence length]

    # Determine the feature map size based on sequence length
    seq_length = current_attn_map.shape[0]
    feature_map_size = int(np.sqrt(seq_length))
    current_attn_map = current_attn_map.view(feature_map_size, feature_map_size)  # Convert to 2D feature map

    # Convert to NumPy array
    current_attn_map = current_attn_map.cpu().numpy()

    # Resize the attention map to match the original image size
    current_attn_map = torch.tensor(current_attn_map).unsqueeze(0).unsqueeze(0)
    current_attn_map = F.interpolate(current_attn_map, size=(original_height, original_width), mode='bilinear', align_corners=False)
    current_attn_map = current_attn_map.squeeze(0).squeeze(0).numpy()

    # Normalize the attention map
    current_attn_map = normalize_scoremap(current_attn_map)

    # Visualization
    axs[i, 0].imshow(original_img)
    axs[i, 0].set_title(f"Original Image - Layer {i+1}")
    axs[i, 0].axis('off')

    axs[i, 1].imshow(original_img)
    axs[i, 1].imshow(current_attn_map, cmap='jet', alpha=0.5)
    axs[i, 1].set_title(f"Attention Map - Layer {i+1}")
    axs[i, 1].axis('off')

    # Add the current layer's attention map to the combined map
    combined_attn_map += current_attn_map

# Normalize the combined attention map
combined_attn_map = normalize_scoremap(combined_attn_map)

# Visualize the combined attention map
axs[num_layers, 0].imshow(original_img)
axs[num_layers, 0].set_title("Original Image - Combined")
axs[num_layers, 0].axis('off')

axs[num_layers, 1].imshow(original_img)
axs[num_layers, 1].imshow(combined_attn_map, cmap='jet', alpha=0.5)
axs[num_layers, 1].set_title("Combined Attention Map")
axs[num_layers, 1].axis('off')

plt.tight_layout()
plt.show()

else:
print("No attention maps were captured!")

Remove hooks

for hook in hooks:
hook.remove()

results：

[ahatamiz]

Hi @mason0417 !

These results look impressive !