File size: 4,255 Bytes
ebcb290 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """
Gradient-weighted Attention Rollout for Vision Transformers (Chefer et al., CVPR 2021).
Calculates class-discriminative attention maps by weighting attention probabilities with backward gradients.
"""
import torch
import numpy as np
class ViTCheferRollout:
def __init__(self, model):
self.model = model
self.attentions = []
self.attention_gradients = []
self.hooks = []
def _register_hooks(self):
self._remove_hooks()
self.attentions = []
self.attention_gradients = []
for name, module in self.model.named_modules():
if "attention.attention" in name and not hasattr(module, "_chefer_hooked"):
def get_attn_hook(module, input, output):
# output in HF ViT self-attention is (context_layer, attn_probs)
if isinstance(output, tuple) and len(output) > 1:
self.attentions.append(output[1])
elif hasattr(output, "attentions"):
self.attentions.append(output.attentions)
def get_attn_grad_hook(module, grad_input, grad_output):
# grad_output[1] is gradient w.r.t attention probabilities
if len(grad_output) > 1 and grad_output[1] is not None:
self.attention_gradients.append(grad_output[1])
h1 = module.register_forward_hook(get_attn_hook)
h2 = module.register_full_backward_hook(get_attn_grad_hook)
self.hooks.extend([h1, h2])
def _remove_hooks(self):
for h in self.hooks:
h.remove()
self.hooks = []
def attribute(self, input_tensor, target_class=0):
"""
Compute Chefer Gradient-Weighted Attention Rollout.
Args:
input_tensor: (1, 3, 384, 384) preprocessed input tensor
target_class: output index (0 for single binary logit)
Returns:
24x24 numpy array of patch attributions
"""
self.model.eval()
self._register_hooks()
device = input_tensor.device
# Enable gradients for backward pass
input_tensor.requires_grad_(True)
self.model.zero_grad()
# Forward pass with output_attentions=True
outputs = self.model(input_tensor, output_attentions=True)
logits = outputs.logits # (1, 1)
# Backward pass on logit
if logits.dim() > 1 and logits.size(1) > target_class:
score = logits[:, target_class]
else:
score = logits.squeeze()
score.backward()
attentions = outputs.attentions # Tuple of 12 (B, heads, 577, 577) tensors
# Cleanup hooks
self._remove_hooks()
if not attentions:
# Fallback if hooks failed to catch: return flat matrix
return np.zeros((24, 24))
seq_len = attentions[0].size(-1) # 577
result = torch.eye(seq_len, device=device)
# Reverse gradient order because backward hook captures in reverse
gradients = self.attention_gradients[::-1] if self.attention_gradients else [None] * len(attentions)
for i, attn in enumerate(attentions):
# attn: (1, heads, 577, 577)
if i < len(gradients) and gradients[i] is not None:
grad = gradients[i]
# Elementwise product: gradient * attention
cam = (grad * attn).clamp(min=0).mean(dim=1) # (1, 577, 577)
else:
# Fallback to plain attention if gradient missing
cam = attn.mean(dim=1)
cam = cam[0] # (577, 577)
cam = cam + torch.eye(seq_len, device=device)
cam = cam / cam.sum(dim=-1, keepdim=True)
result = torch.matmul(cam, result)
# Extract CLS token -> patch tokens
patch_attn = result[0, 1:].reshape(24, 24).detach().cpu().numpy()
return patch_attn
def chefer_gradient_rollout(wrapper, input_tensor) -> np.ndarray:
"""Convenience function for running Chefer rollout on ViTWrapper."""
explainer = ViTCheferRollout(wrapper.model)
return explainer.attribute(input_tensor, target_class=0)
|