V1.6.0: Refactor the core dependency module

V1.6.0: Refactor the core dependency module

Author: VainF

README.md

@@ -2,16 +2,16 @@
 <div align="center">
 <img src="https://user-images.githubusercontent.com/18592211/232830417-0b21a874-516e-4420-8984-4de414a35085.png" width="400px"></img>
 <h2></h2>
-<h3>Towards Any Structural Pruning<h3>
-<img src="assets/intro.png" width="50%">
+<img src="https://github.com/user-attachments/assets/50b03774-7345-4eb6-bf28-209195d354b0" width="40%">
+<h2></h2>
 </div>
 
 <p align="center">
   <a href="https://github.com/VainF/Torch-Pruning/actions"><img src="https://img.shields.io/badge/tests-passing-9c27b0.svg" alt="Test Status"></a>
   <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-1.x %20%7C%202.x-673ab7.svg" alt="Tested PyTorch Versions"></a>
   <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-4caf50.svg" alt="License"></a>
   <a href="https://pepy.tech/project/Torch-Pruning"><img src="https://static.pepy.tech/badge/Torch-Pruning?color=2196f3" alt="Downloads"></a>
-  <a href="https://github.com/VainF/Torch-Pruning/releases/latest"><img src="https://img.shields.io/badge/Latest%20Version-1.5.3-3f51b5.svg" alt="Latest Version"></a>
+  <a href="https://github.com/VainF/Torch-Pruning/releases/latest"><img src="https://img.shields.io/badge/Latest%20Version-1.6.0-3f51b5.svg" alt="Latest Version"></a>
   <a href="https://colab.research.google.com/drive/1TRvELQDNj9PwM-EERWbF3IQOyxZeDepp?usp=sharing">
   <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
 </a>

examples/timm_models/prune_timm_models.py

@@ -9,6 +9,8 @@
 from timm.models.vision_transformer import Attention
 import torch_pruning as tp
 import argparse
+from typing import Optional, Type
+from timm.models.vision_transformer import maybe_add_mask
 
 parser = argparse.ArgumentParser(description='Prune timm models')
 parser.add_argument('--model', default=None, type=str, help='model name')
@@ -18,9 +20,11 @@
 parser.add_argument('--list_models', default=False, action='store_true', help='list all models in timm')
 args = parser.parse_args()
 
-
-def forward(self, x):
-    """https://github.com/huggingface/pytorch-image-models/blob/054c763fcaa7d241564439ae05fbe919ed85e614/timm/models/vision_transformer.py#L79"""
+def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
     B, N, C = x.shape
     qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
     q, k, v = qkv.unbind(0)
@@ -29,21 +33,23 @@ def forward(self, x):
     if self.fused_attn:
         x = F.scaled_dot_product_attention(
             q, k, v,
-            dropout_p=self.attn_drop.p,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_drop.p if self.training else 0.,
         )
     else:
         q = q * self.scale
         attn = q @ k.transpose(-2, -1)
+        attn = maybe_add_mask(attn, attn_mask)
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
         x = attn @ v
 
-    x = x.transpose(1, 2).reshape(B, N, -1) # original implementation: x = x.transpose(1, 2).reshape(B, N, C)
+    x = x.transpose(1, 2).reshape(B, N, -1)
+    x = self.norm(x)
     x = self.proj(x)
     x = self.proj_drop(x)
     return x
 
-
 def main():
     timm_models = timm.list_models()
     if args.list_models:

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="torch-pruning",
-    version="v1.5.3",
+    version="v1.6.0",
     author="Gongfan Fang",
     author_email="gongfan@u.nus.edu",
     description="Towards Any Structural Pruning",

torch_pruning/_helpers.py

@@ -42,91 +42,6 @@ def is_scalar(x):
         return False
     return False
 
-
-class _FlattenIndexMapping(object):
-    def __init__(self, stride=1, reverse=False):
-        self._stride = stride
-        self.reverse = reverse
-
-    def __call__(self, idxs: _HybridIndex):
-        new_idxs = []
-        
-        if self.reverse == True:
-            for i in idxs:
-                new_idxs.append( _HybridIndex( idx = (i.idx // self._stride), root_idx=i.root_idx ) )
-            new_idxs = list(set(new_idxs))
-        else:
-            for i in idxs:
-                new_idxs.extend(
-                    [ _HybridIndex(idx=k, root_idx=i.root_idx) for k in range(i.idx * self._stride, (i.idx + 1) * self._stride) ]  
-                )
-        return new_idxs
-
-
-class _ConcatIndexMapping(object):
-    def __init__(self, offset, reverse=False):
-        self.offset = offset
-        self.reverse = reverse
-
-    def __call__(self, idxs: _HybridIndex):
-        if self.reverse == True:
-            new_idxs = [
-                _HybridIndex(idx = i.idx - self.offset[0], root_idx=i.root_idx )
-                for i in idxs
-                if (i.idx >= self.offset[0] and i.idx < self.offset[1])
-            ]
-        else:
-            new_idxs = [ _HybridIndex(idx=i.idx + self.offset[0], root_idx=i.root_idx) for i in idxs]
-        return new_idxs
-
-class _GQAIndexMapping(object):
-    def __init__(self, repeat, head_dim, reverse=False):
-        self.repeat = repeat
-        self.reverse = reverse
-        self.head_dim = head_dim
-
-    def __call__(self, idxs: _HybridIndex):
-        head_dim = self.head_dim
-        repeat = self.repeat
-        if self.reverse == True: 
-            new_idxs = [ _HybridIndex(idx=( i.idx - i.idx // (head_dim * repeat) * head_dim * (repeat - 1) - i.idx//head_dim%repeat * head_dim ), root_idx=None) for i in idxs ]
-        else:
-            new_idxs = []
-           
-        return new_idxs
-
-class _SliceIndexMapping(object):
-    def __init__(self, dim, start, step, end, reverse=False):
-        self.start = start
-        self.step = step
-        self.end = end
-        self.reverse = reverse
-        self.dim = dim
-    
-    def __call__(self, idxs: _HybridIndex):
-        
-        if self.reverse == True:
-            new_idxs = [ _HybridIndex(idx=i.idx * self.step + self.start, root_idx=i.root_idx) for i in idxs]
-        else:
-            new_idxs = [ _HybridIndex(idx=(i.idx - self.start) // self.step, root_idx=i.root_idx) for i in idxs if (i.idx >= self.start and i.idx < self.end and (i.idx-self.start)%self.step==0) ]
-        return new_idxs
-
-class _SplitIndexMapping(object):
-    def __init__(self, offset, reverse=False):
-        self.offset = offset
-        self.reverse = reverse
-
-    def __call__(self, idxs: _HybridIndex):
-        if self.reverse == True:
-            new_idxs = [ _HybridIndex(idx=i.idx + self.offset[0], root_idx=i.root_idx) for i in idxs]
-        else:
-            new_idxs = [
-                _HybridIndex(idx = i.idx - self.offset[0], root_idx=i.root_idx)
-                for i in idxs
-                if (i.idx >= self.offset[0] and i.idx < self.offset[1])
-            ]
-        return new_idxs
-
 class ScalarSum:
     def __init__(self):
         self._results = {}