Spaces:
Running
Running
File size: 3,560 Bytes
c64fb9f |
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 |
import time
import numpy as np
import torch
from PIL import Image
import glob
import sys
import argparse
import datetime
import json
from pathlib import Path
from llava.hook import HookManager
def init_hookmanager(module):
module.hook_manager = HookManager()
class MaskHookLogger(object):
def __init__(self, model, device):
self.current_layer = 0
self.device = device
self.attns = []
self.projected_attns = []
self.image_embed_range = []
self.index = []
self.model = model
@torch.no_grad()
def compute_attentions(self, ret):
assert len(self.image_embed_range) > 0
st, ed = self.image_embed_range[-1]
image_attention = ret[:,:,-1,st:ed].detach()
image_attention = image_attention.mean(dim = 1)
self.attns.append(image_attention) # [b, k]
return ret
@torch.no_grad()
def compute_projected_attentions(self, ret):
assert len(self.image_embed_range) > 0
st, ed = self.image_embed_range[-1]
image_attention = ret[:,-1,st:ed].detach() # [b, k, d]
self.projected_attns.append(image_attention) # [b, k, d]
return ret
@torch.no_grad()
def compute_attentions_withsoftmax(self, ret):
assert len(self.image_embed_range) > 0
st, ed = self.image_embed_range[-1]
image_attention = ret[:,:,-1,st:ed].detach()
image_attention = image_attention.softmax(dim = -1)
image_attention = image_attention.mean(dim = 1)
self.attns.append(image_attention) # [b, k]
return ret
@torch.no_grad()
def compute_logits_index(self, ret):
next_token_logits = ret[:, -1, :]
index = next_token_logits.argmax(dim=-1)
self.index.append(index.item())
return ret
@torch.no_grad()
def finalize(self):
attns = torch.cat(self.attns, dim = 0).to(self.device)
return attns
@torch.no_grad()
def finalize_projected_attn(self, norm_weight, proj):
assert len(self.index) == len(self.projected_attns)
mask = []
for i in range(-4,-2):
index = self.index[i]
attns = self.projected_attns[i].to(self.device) # 1,k,d
input_dtype = attns.dtype
attns_var = attns.to(torch.float32).sum(dim = 1).pow(2).mean(-1, keepdim=True)# 1,d
attns_var = attns_var.unsqueeze(1)# 1,1,d
normalized_attns = attns * torch.rsqrt(attns_var + 1e-6) # 1,k,d
normalized_attns = norm_weight.to(normalized_attns.device) * normalized_attns.to(input_dtype) # 1,k,d
logits = proj(normalized_attns)
max_logits = logits[0,:,index] # k
mask.append(max_logits)
mask = torch.stack(mask, dim = 0)
return mask.mean(dim = 0)
def reinit(self):
self.attns = []
self.projected_attns = []
self.image_embed_range = []
self.index = []
torch.cuda.empty_cache()
def log_image_embeds_range(self, ret):
self.image_embed_range.append(ret[0][0])
return ret
def hook_logger(model, device, layer_index = 20):
"""Hooks a projected residual stream logger to the model."""
init_hookmanager(model.model.layers[layer_index].self_attn)
prs = MaskHookLogger(model, device)
model.model.layers[layer_index].self_attn.hook_manager.register('after_attn_mask',
prs.compute_attentions_withsoftmax)
model.hooklogger = prs
return prs
|