prunevid (#389)

Author: chengtao-lv

configs/sparsification/methods/PruneVid/prunevid.yml

@@ -0,0 +1,28 @@
+base:
+    seed: &seed 42
+model:
+    type: Llava OneVision
+    path: model path
+    torch_dtype: auto
+eval:
+    eval_pos: [pretrain, transformed]
+    type: vqa
+    name: [mme]
+    download: False
+    path: MME dataset path
+    bs: 1
+    inference_per_block: False
+sparse:
+    method: TokenReduction
+    special:
+        method: PruneVid
+        lora_alpha: 14
+        selected_layers: 10
+        alphas: 0.4
+        taus: 0.8
+        temporal_segment_ratios: 0.25
+        cluster_ratios: 0.5
+save:
+    save_trans: False
+    save_fake: False
+    save_path: /path/to/save/

llmc/compression/token_reduction/__init__.py

@@ -2,6 +2,7 @@
 from .dycoke import DyCoke
 from .fastervlm import FasterVLM
 from .fastv import FastV
+from .prunevid import PruneVid
 from .pyramiddrop import PyramidDrop
 from .sparsevlm import SparseVLM
 from .tome import ToMe

llmc/compression/token_reduction/dycoke.py

@@ -14,7 +14,7 @@
 from llmc.utils.registry_factory import TOKEN_REDUCTION_REGISTRY
 
 from .token_reduction_module import TokenReductionModule
-from .utils import prefill_wrapper
+from .utils import add_post_hook_to_get_2dPool
 
 
 def dycole_ttm(image_feature, pruning_paras):
@@ -102,16 +102,6 @@ def dycole_ttm(image_feature, pruning_paras):
     return combined_tokens
 
 
-def add_dycole_ttm_to_get_2dPool(model, post_hook_fn, pruning_paras):
-    original_fn = model.get_2dPool
-
-    def wrapped_fn(*args, **kwargs):
-        result = original_fn(*args, **kwargs)
-        return post_hook_fn(result, pruning_paras)
-
-    model.get_2dPool = wrapped_fn
-
-
 @TOKEN_REDUCTION_REGISTRY.register('DyCoke')
 class DyCoke(TokenReductionModule):
     def __init__(self, config, model, blocks):
@@ -127,6 +117,6 @@ def add_sparse_config(self):
     def register_reduction_modules(self):
 
         if isinstance(self.model.model, LlavaMetaForCausalLM):
-            add_dycole_ttm_to_get_2dPool(
+            add_post_hook_to_get_2dPool(
                 self.model.model, dycole_ttm, self.model.model.pruning_paras
             )