I want to visualize the attention map.
We primarily visualize attention maps using values after attn and softmax calculations.
self.qk = nn.Conv2d(dim, all_head_dim * 2, 1, bias=False)
self.v = nn.Conv2d(dim, all_head_dim, 1, bias=False)
self.proj = nn.Conv2d(all_head_dim, dim, 1, bias=False)
self.pe = Conv(all_head_dim, dim, 5, 1, 2, g=dim, act=False)
self.last_attention = None
self.save_dir = Path(save_dir)
self.save_dir.mkdir(exist_ok=True, parents=True)
self.forward_count = 0 # 保存用カウンタ
def forward(self, x):
print("AAttn forward called")
B, C, H, W = x.shape
print(x.shape)
N = H * W
# Q, K, Vの計算
qk = self.qk(x).flatten(2).transpose(1, 2) # [B, N, C*2]
print(qk.shape)
v = self.v(x)
print(v.shape)
pp = self.pe(v)
print(pp.shape)
v = v.flatten(2).transpose(1, 2)
print(v.shape)
if self.area > 1:
qk = qk.reshape(B * self.area, N // self.area, C * 2)
print("qk:", qk.shape)
v = v.reshape(B * self.area, N // self.area, C)
print("v:", v.shape)
B, N, _ = qk.shape
q, k = qk.split([C, C], dim=2)
print("q:", q.shape)
print("k:", k.shape)
USE_FLASH_ATTN = False
if x.is_cuda and USE_FLASH_ATTN:
q = q.view(B, N, self.num_heads, self.head_dim)
k = k.view(B, N, self.num_heads, self.head_dim)
v = v.view(B, N, self.num_heads, self.head_dim)
x = flash_attn_func(
q.contiguous().half(),
k.contiguous().half(),
v.contiguous().half()
).to(q.dtype)
else:
q = q.transpose(1, 2).view(B, self.num_heads, self.head_dim, N)
k = k.transpose(1, 2).view(B, self.num_heads, self.head_dim, N)
v = v.transpose(1, 2).view(B, self.num_heads, self.head_dim, N)
print("qnf:", q.shape)
print("knf:", k.shape)
print("vnf:", v.shape)
# アテンション計算
attn = (q.transpose(-2, -1) @ k) * (self.head_dim ** -0.5)
max_attn = attn.max(dim=-1, keepdim=True).values
exp_attn = torch.exp(attn - max_attn)
attn_normalized = exp_attn / exp_attn.sum(dim=-1, keepdim=True)
print(attn_normalized.shape)
# アテンションマップを保存
self.last_attention = attn_normalized.detach().cpu()
# ★★★ ここで実際に保存 ★★★
self.save_attention_as_image(attn_normalized, H, W)
# 続きの処理
x = (v @ attn_normalized.transpose(-2, -1))
print("xnfb:", x.shape)
x = x.permute(0, 3, 1, 2)
print("xnfa:", x.shape)
if self.area > 1:
x = x.reshape(B // self.area, N * self.area, C)
print("x2:", x.shape)
B, N, _ = x.shape
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
print("xl:", x.shape)
z = self.proj(x + pp)
print("z:", z.shape)
return self.proj(x + pp)
# H/l
def save_attention_as_image(self, attn, H, W):
try:
area = self.area
B_area, num_heads, N, _ = attn.shape
assert B_area % area == 0
H_area = H // area
canvas = torch.zeros(H, W)
for a in range(area):
idx = a # batch0
attn_map = attn[idx, 0] # [N, N]
mean_attn = attn_map.mean(dim=0) # [N]
# H分割なのでこちらが正解
area_attn_2d = mean_attn.view(H_area, W)
# 正規化
area_attn_2d -= area_attn_2d.min()
area_attn_2d /= (area_attn_2d.max() + 1e-8)
h_start = a * H_area
h_end = (a + 1) * H_area
canvas[h_start:h_end, :] = area_attn_2d
canvas = canvas.cpu().numpy()
plt.figure(figsize=(8, 6))
plt.imshow(canvas, cmap="hot")
plt.colorbar()
plt.title("Area Attention Map (H-split)")
plt.axis("off")
filename = self.save_dir / f"attention_stitched_{self.forward_count:04d}.png"
plt.savefig(filename, dpi=150, bbox_inches="tight")
plt.close()
print(f"✓ Stitched attention map saved: {filename}")
self.forward_count += 1
except Exception as e:
print(f"Error saving stitched attention map: {e}")
I want to visualize the attention map.
Is the following visualization method correct?
(FLASH_Visualizing attention maps when no attention is applied.)
We primarily visualize attention maps using values after attn and softmax calculations.
`
class AAttn(nn.Module):
def init(self, dim, num_heads, area=1, save_dir="attention_maps"):
super().init()
self.area = area
self.num_heads = num_heads
self.head_dim = head_dim = dim // num_heads
all_head_dim = head_dim * self.num_heads
`