File size: 3,156 Bytes
a20151e | 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 | import torch
import numpy as np
import csv
from transformers import ViTModel
# Load ViT
model = ViTModel.from_pretrained("google/vit-base-patch16-224")
hidden_size = model.config.hidden_size
num_heads = model.config.num_attention_heads
head_dim = hidden_size // num_heads
# Norm calculator
def compute_norms(A: torch.Tensor):
return {
"1": torch.norm(A, p=1).item(),
"F": torch.norm(A, p="fro").item(),
"*": torch.linalg.svdvals(A).sum().item(),
"2,1": torch.norm(A, dim=0, p=2).sum().item(),
"2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item()
}
norm_names = ["1","F","*","2,1","2,1,T"]
subcols = ["Q","K","Q/K","V","O","V/O"]
outfile = "vit_qkvo_norms.csv"
with open(outfile, "w", newline="") as f:
writer = csv.writer(f)
for layer_idx, layer in enumerate(model.encoder.layer, start=1):
# Header rows
header1 = [f"Layer {layer_idx}"]
for n in norm_names:
header1.extend([n,"","","","",""])
writer.writerow(header1)
header2 = [""]
for _ in norm_names:
header2.extend(subcols)
writer.writerow(header2)
rows = []
# Extract weights + biases
W_q = torch.cat([layer.attention.attention.query.weight.detach(),
layer.attention.attention.query.bias.detach().unsqueeze(1)], dim=1)
W_k = torch.cat([layer.attention.attention.key.weight.detach(),
layer.attention.attention.key.bias.detach().unsqueeze(1)], dim=1)
W_v = torch.cat([layer.attention.attention.value.weight.detach(),
layer.attention.attention.value.bias.detach().unsqueeze(1)], dim=1)
W_o = torch.cat([layer.attention.output.dense.weight.detach(),
layer.attention.output.dense.bias.detach().unsqueeze(1)], dim=1)
# Reshape into heads
W_q_heads = W_q.view(num_heads, head_dim, -1)
W_k_heads = W_k.view(num_heads, head_dim, -1)
W_v_heads = W_v.view(num_heads, head_dim, -1)
W_o_heads = W_o.view(num_heads, head_dim, -1)
for h in range(num_heads):
row = [f"Head {h+1}"]
for norm in norm_names:
nq = compute_norms(W_q_heads[h])[norm]
nk = compute_norms(W_k_heads[h])[norm]
nv = compute_norms(W_v_heads[h])[norm]
no = compute_norms(W_o_heads[h])[norm]
qk_ratio = nq/(nk+1e-12)
vo_ratio = nv/(no+1e-12)
row.extend([
round(nq,4),
round(nk,4),
round(qk_ratio,4),
round(nv,4),
round(no,4),
round(vo_ratio,4)
])
writer.writerow(row)
rows.append(row[1:])
# Mean & Std rows
arr = np.array(rows, dtype=float)
mean = np.round(arr.mean(axis=0),4)
std = np.round(arr.std(axis=0),4)
writer.writerow(["Mean"] + mean.tolist())
writer.writerow(["Std"] + std.tolist())
writer.writerow([])
print(f"✅ Saved CSV: {outfile}")
|