add vision zip

Author: chengtao-lv

configs/sparsification/methods/VisionZip/visionzip.yml

@@ -0,0 +1,24 @@
+base:
+    seed: &seed 42
+model:
+    type: Llava
+    path: model path
+    torch_dtype: auto
+eval:
+    eval_pos: [transformed]
+    type: vqa
+    name: [mme]
+    download: False
+    path: MME dataset path
+    bs: 1
+    inference_per_block: False
+sparse:
+    method: TokenReduction
+    special:
+        method: VisionZip
+        dominant: 191
+        contextual: 30
+save:
+    save_trans: False
+    save_fake: False
+    save_path: /path/to/save/

llmc/compression/token_reduction/__init__.py

@@ -3,3 +3,4 @@
 from .fastv import FastV
 from .sparsevlm import SparseVLM
 from .tome import ToMe
+from .visionzip import VisionZip

llmc/compression/token_reduction/sparsevlm.py

@@ -1,12 +1,7 @@
 import functools
-import math
-import types
-from typing import Callable, Optional, Tuple
 
 import einops as ein
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 
 from llmc.utils.registry_factory import TOKEN_REDUCTION_REGISTRY
 

llmc/compression/token_reduction/utils.py

@@ -0,0 +1,70 @@
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers.models.clip.modeling_clip import CLIPEncoderLayer
+
+
+def parse_r(num_layers: int, r: Union[List[int], Tuple[int, float], int]) -> List[int]:
+    """Copy from the TOME. https://github.com/facebookresearch/ToMe.
+
+    Process a constant r or r schedule into a list for use internally.
+
+    r can take the following forms:
+     - int: A constant number of tokens per layer.
+     - Tuple[int, float]: A pair of r, inflection.
+       Inflection describes there the the reduction / layer should trend
+       upward (+1), downward (-1), or stay constant (0). A value of (r, 0)
+       is as providing a constant r. (r, -1) is what we describe in the paper
+       as "decreasing schedule". Any value between -1 and +1 is accepted.
+     - List[int]: A specific number of tokens per layer. For extreme granularity.
+    """
+    inflect = 0
+    if isinstance(r, list):
+        if len(r) < num_layers:
+            r = r + [0] * (num_layers - len(r))
+        return list(r)
+    elif isinstance(r, tuple):
+        r, inflect = r
+
+    min_val = int(r * (1.0 - inflect))
+    max_val = 2 * r - min_val
+    step = (max_val - min_val) / (num_layers - 1)
+
+    return [int(min_val + step * i) for i in range(num_layers)]
+
+
+def make_tome_class(transformer_class):
+    class VisionZipTransformer(transformer_class):
+        """
+        Modifications:
+        - Initialize r, token size, and token sources.
+        """
+
+        def forward(self, *args, **kwdargs) -> torch.Tensor:
+            self._info['r'] = parse_r(len(self.vision_model.encoder.layers), self.r)
+            # self._info["r"] = self.r
+
+            self._info['size'] = None
+            self._info['source'] = None
+
+            return super().forward(*args, **kwdargs)
+
+    return VisionZipTransformer
+
+
+def apply_info(model, dominant_num, contextual_num):
+
+    VisionZipTransformer = make_tome_class(model.__class__)
+
+    model.__class__ = VisionZipTransformer
+    model.r = [0 for i in range(22)] + [1] + [0]
+
+    model._info = {
+        'r': [model.r],
+        'dominant': dominant_num,
+        'contextual': contextual_num,
+    }
+    for module in model.modules():
+        if isinstance(module, CLIPEncoderLayer):
+            module.self_attn.k_proj._info = model._info