|
| 1 | +import functools |
| 2 | +import torch |
| 3 | +from llmc.utils.registry_factory import TOKEN_REDUCTION_REGISTRY |
| 4 | +from .token_reduction_module import TokenReductionModule |
| 5 | + |
| 6 | + |
| 7 | +@TOKEN_REDUCTION_REGISTRY.register('VisPruner') |
| 8 | +class VisPruner(TokenReductionModule): |
| 9 | + def __init__(self, config, model, blocks): |
| 10 | + super().__init__(config, model, blocks) |
| 11 | + self.add_sparse_config() |
| 12 | + self.register_reduction_modules() |
| 13 | + |
| 14 | + def add_sparse_config(self): |
| 15 | + self.special_config['select_layer'] = self.model.pruning_config.get('select_layer', -1) |
| 16 | + self.special_config['select_feature'] = self.model.pruning_config.get('select_feature', None) |
| 17 | + |
| 18 | + self.pruning_paras = self.special_config |
| 19 | + |
| 20 | + def register_reduction_modules(self): |
| 21 | + |
| 22 | + def update_output_attentions_hook(module, args, kwargs): |
| 23 | + kwargs['output_attentions'] = True |
| 24 | + |
| 25 | + def store_attention_hook(module, inps, outs, pruning_paras): |
| 26 | + image_attentions = outs.attentions[pruning_paras['select_layer']] |
| 27 | + if pruning_paras['select_feature'] == 'patch': |
| 28 | + image_attentions = image_attentions[:, :, 0, 1:] |
| 29 | + elif pruning_paras['select_feature'] == 'cls_patch': |
| 30 | + image_attentions = image_attentions |
| 31 | + raise ValueError(f'Unexpected select feature: {self.select_feature}') |
| 32 | + |
| 33 | + pruning_paras['image_attentions'] = image_attentions.to(inps[0].dtype) |
| 34 | + |
| 35 | + def get_index_masks_hook(module, args, pruning_paras): |
| 36 | + image_features = args[0] |
| 37 | + image_attentions = pruning_paras['image_attentions'] |
| 38 | + |
| 39 | + B, N, C = image_features.shape |
| 40 | + device = image_features.device |
| 41 | + index_masks = torch.ones(B, N, dtype=torch.bool, device=device) |
| 42 | + |
| 43 | + visual_token_num = round( |
| 44 | + self.special_config['vision_token_length'] * (1 - self.special_config['prune_ratio']) |
| 45 | + ) # T |
| 46 | + important_ratio = self.pruning_paras['important_ratio'] # r |
| 47 | + important_token_num = int(visual_token_num * important_ratio) # T_imp = T * r |
| 48 | + diverse_token_num = visual_token_num - important_token_num # T_div = T * (1 - r) |
| 49 | + |
| 50 | + # [VisPruner] Select important tokens using attention scores |
| 51 | + image_attentions = image_attentions.mean(dim=1) # (B, N) |
| 52 | + token_indices = image_attentions.argsort(dim=-1, descending=True) # (B, N) |
| 53 | + important_indices = token_indices[:, :important_token_num] # (B, T_imp) |
| 54 | + residual_indices = token_indices[:, important_token_num:] # (B, N - T_imp) |
| 55 | + |
| 56 | + # [VisPruner] Remove duplicate tokens by iterative matching and pruning |
| 57 | + image_normalized = image_features / image_features.norm(dim=-1, keepdim=True) # (B, N, C) |
| 58 | + while diverse_token_num > 0: |
| 59 | + R = residual_indices.shape[1] |
| 60 | + r = min(8, R - diverse_token_num) |
| 61 | + if r <= 0: |
| 62 | + break |
| 63 | + |
| 64 | + residual_tokens = image_normalized[torch.arange(B).unsqueeze(-1).expand(-1, R), residual_indices] # (B, R, C) |
| 65 | + a, b = residual_tokens[..., ::2, :], residual_tokens[..., 1::2, :] # (B, R // 2, C) |
| 66 | + scores = a @ b.transpose(-1, -2) # (B, R // 2, R // 2) |
| 67 | + scores = scores.max(dim=-1).values # (B, R // 2) |
| 68 | + |
| 69 | + distinct_indices = scores.argsort(dim=-1, descending=True)[:, r:] # (B, R // 2 - r) |
| 70 | + residual_indices = torch.cat([ |
| 71 | + residual_indices[..., ::2][torch.arange(B).unsqueeze(-1).expand(-1, R // 2 - r), distinct_indices], |
| 72 | + residual_indices[..., 1::2] |
| 73 | + ], dim=-1) # (B, R - r) |
| 74 | + |
| 75 | + if diverse_token_num > 0: |
| 76 | + selected_indices = torch.cat([important_indices, residual_indices], dim=-1) # (B, T) |
| 77 | + else: |
| 78 | + selected_indices = important_indices # (B, T) |
| 79 | + index_masks = torch.zeros(B, N, dtype=torch.bool, device=device) |
| 80 | + index_masks.scatter_(1, selected_indices, True) |
| 81 | + |
| 82 | + pruning_paras['index_masks'] = index_masks |
| 83 | + |
| 84 | + def prune_hook(module, inputs, outputs, pruning_paras): |
| 85 | + image_features = outputs |
| 86 | + index_masks = pruning_paras['index_masks'] |
| 87 | + return image_features[index_masks].unsqueeze(0) |
| 88 | + |
| 89 | + self.model.vision_model.vision_tower.register_forward_pre_hook( |
| 90 | + update_output_attentions_hook, |
| 91 | + with_kwargs=True |
| 92 | + ) |
| 93 | + |
| 94 | + self.model.vision_model.vision_tower.register_forward_hook( |
| 95 | + functools.partial(store_attention_hook, pruning_paras=self.pruning_paras), |
| 96 | + ) |
| 97 | + |
| 98 | + self.model.vision_projector.register_forward_pre_hook( |
| 99 | + functools.partial(get_index_masks_hook, pruning_paras=self.pruning_paras), |
| 100 | + ) |
| 101 | + |
| 102 | + self.model.vision_projector.register_forward_hook( |
| 103 | + functools.partial(prune_hook, pruning_paras=self.pruning_paras), |
| 104 | + ) |
| 105 | + |
| 106 | + |
0 commit comments