Repository: VainF/Torch-Pruning (language: python)

Issue:
V1.2.5



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 8c54fe2..250657f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -24,4 +24,7 @@ torch_pruning_bak
 .ipynb_checkpoints
 .idea/
 .pytest_cache
-benchmarks/run
\ No newline at end of file
+benchmarks/run
+output
+acc.png
+pretrained
\ No newline at end of file
diff --git a/README.md b/README.md
index 0ec3893..51d31cc 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@
   <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-1.8 %20%7C%201.12 %20%7C%202.0-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.2.4-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.2.5-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>
@@ -23,7 +23,7 @@
 
 Torch-Pruning (TP) is a library for structural pruning with the following features:
 
-* **General-purpose Pruning Toolkit:** TP enables structural pruning for a wide range of deep neural networks, including [Large Language Models (LLMs)](https://github.com/horseee/LLM-Pruner), [Diffusion Models](https://github.com/VainF/Diff-Pruning), [Yolov7](examples/yolov7/), [yolov8](examples/yolov8/), [Vision Transformers](examples/hf_transformers/), [Swin Transformers](examples/hf_transformers), [BERT](examples/hf_transformers), FasterRCNN, SSD, ResNe(X)t, ConvNext, DenseNet, ConvNext, RegNet, DeepLab, etc. Different from [torch.nn.utils.prune](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html) that zeroizes parameters through masking, Torch-Pruning deploys a (non-deep) graph algorithm called **[DepGraph](https://openaccess.thecvf.com/content/CVPR2023/html/Fang_DepGraph_Towards_Any_Structural_Pruning_CVPR_2023_paper.html)** to remove parameters physically. 
+* **General-purpose Pruning Toolkit:** TP enables structural pruning for a wide range of deep neural networks, including [Large Language Models (LLMs)](https://github.com/horseee/LLM-Pruner), [Diffusion Models](https://github.com/VainF/Diff-Pruning), [Yolov7](examples/yolov7/), [yolov8](examples/yolov8/), [Vision Transformers](examples/transformers/), [Swin Transformers](examples/transformers#swin-transformers-from-hf-transformers), [BERT](examples/transformers#bert-from-hf-transformers), FasterRCNN, SSD, ResNe(X)t, ConvNext, DenseNet, ConvNext, RegNet, DeepLab, etc. Different from [torch.nn.utils.prune](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html) that zeroizes parameters through masking, Torch-Pruning deploys an algorithm called **[DepGraph](https://openaccess.thecvf.com/content/CVPR2023/html/Fang_DepGraph_Towards_Any_Structural_Pruning_CVPR_2023_paper.html)** to remove parameters physically. Check our [Paper List](https://github.com/VainF/Torch-Pruning/wiki/0.-Paper-List) for more details.
 * **[Examples](examples)**: Play around with off-the-shelf models from Huggingface Transformers, Timm, Torchvision, Yolo, etc.  
 * **[Benchmark](benchmarks)**: Reproduce the our results in the DepGraph paper.
 
@@ -33,14 +33,14 @@ For more technical details, please refer to our CVPR'23 paper:
 > *[Learning and Vision Lab](http://lv-nus.org/), National University of Singapore*
   
 ### Update:
-* 2023.09.06 [Pruning & Finetuning Examples for Vision Transformers, Swin Transformers, Bert](examples/transformers/).
-* 2023.07.19 :rocket: Support LLaMA, LLaMA-2, Vicuna, Baichuan, Bloom in [LLM-Pruner](https://github.com/horseee/LLM-Pruner)
-* 2023.05.20 :rocket: [**LLM-Pruner: On the Structural Pruning of Large Language Models**](https://github.com/horseee/LLM-Pruner)  [*[arXiv]*](https://arxiv.org/abs/2305.11627)
-* 2023.05.19 [Structural Pruning for Diffusion Models](https://github.com/VainF/Diff-Pruning) [*[arXiv]*](https://arxiv.org/abs/2305.10924)
-* 2023.04.15 [Pruning and Post-training for YOLOv7 / YOLOv8](benchmarks/examples)
+- [x] 2023.09.06 [Pruning & Finetuning Examples for Vision Transformers, Swin Transformers, Bert](examples/transformers/).
+- [x] 2023.07.19 :rocket: Support LLaMA, LLaMA-2, Vicuna, Baichuan, Bloom in [LLM-Pruner](https://github.com/horseee/LLM-Pruner)
+- [x] 2023.05.20 :rocket: [**LLM-Pruner: On the Structural Pruning of Large Language Models**](https://github.com/horseee/LLM-Pruner)  [*[arXiv]*](https://arxiv.org/abs/2305.11627)
+- [x] 2023.05.19 [Structural Pruning for Diffusion Models](https://github.com/VainF/Diff-Pruning) [*[arXiv]*](https://arxiv.org/abs/2305.10924)
+- [x] 2023.04.15 [Pruning and Post-training for YOLOv7 / YOLOv8](benchmarks/examples)
 
 ### **Features:**
-- [x] High-level Pruners: [MetaPruner](torch_pruning/pruner/algorithms/metapruner.py), [MagnitudePruner](https://arxiv.org/abs/1608.08710), [BNScalePruner](https://arxiv.org/abs/1708.06519), [GroupNormPruner](https://arxiv.org/abs/2301.12900), [GrowingRegPruner](https://arxiv.org/abs/2012.09243), RandomPruner, etc. A paper list is available on our [wiki page](https://github.com/VainF/Torch-Pruning/wiki/Paper-List).
+- [x] High-level Pruners: [MetaPruner](torch_pruning/pruner/algorithms/metapruner.py), [MagnitudePruner](https://arxiv.org/abs/1608.08710), [BNScalePruner](https://arxiv.org/abs/1708.06519), [GroupNormPruner](https://arxiv.org/abs/2301.12900), [GrowingRegPruner](https://arxiv.org/abs/2012.09243), RandomPruner, etc. A paper list is available on our [wiki page](https://github.com/VainF/Torch-Pruning/wiki/0.-Paper-List).
 - [x] Dependency Graph for automatic structural pruning
 - [x] [Low-level pruning functions](torch_pruning/pruner/function.py)
 - [x] Supported Importance Criteria: L-p Norm, Taylor, Random, BNScaling, etc.
@@ -51,7 +51,7 @@ For more technical details, please refer to our CVPR'23 paper:
 ### **Contact Us:**
 Please do not hesitate to open a [discussion](https://github.com/VainF/Torch-Pruning/discussions) or [issue](https://github.com/VainF/Torch-Pruning/issues) if you encounter any problems with the library or the paper.   
 Or Join our Discord or WeChat group for a chat:
-  * Discord: [link](https://discord.gg/2CXnf5bN)
+  * Discord: [link](https://discord.gg/hjX24XA9)
   * WeChat (Group size exceeded 200): [QR Code](https://github.com/VainF/Torch-Pruning/assets/18592211/6c80e758-7692-4dad-b6aa-1e1877e72bf7)
 
 ## Installation
@@ -92,7 +92,7 @@ model = resnet18(pretrained=True).eval()
 # 1. Build dependency graph for resnet18
 DG = tp.DependencyGraph().build_dependency(model, example_inputs=torch.randn(1,3,224,224))
 
-# 2. Group all coupled layers
+# 2. Group coupled layers for model.conv1
 group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
 
 # 3. Prune grouped layers altogether
@@ -154,34 +154,36 @@ model = resnet18(pretrained=True)
 example_inputs = torch.randn(1, 3, 224, 224)
 
 # 1. Importance criterion
-imp = tp.importance.GroupTaylorImportance() # or GroupNormPruner, GroupHessianImportance, etc.
+imp = tp.importance.GroupTaylorImportance() # or GroupNormImportance(p=2), GroupHessianImportance(), etc.
 
 # 2. Initialize a pruner with the model and the importance criterion
 ignored_layers = []
 for m in model.modules():
     if isinstance(m, torch.nn.Linear) and m.out_features == 1000:
         ignored_layers.append(m) # DO NOT prune the final classifier!
+
 pruner = tp.pruner.MetaPruner( # We can always choose MetaPruner if sparse training is not required.
     model,
     example_inputs,
     importance=imp,
     ch_sparsity=0.5, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
+    # ch_sparsity_dict = {model.conv1: 0.2, model.layer2: 0.8}, # customized sparsity for layers or blocks
     ignored_layers=ignored_layers,
 )
 
 # 3. Prune & finetune the model
 base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
-for i in range(iterative_steps):
-    if isinstance(imp, tp.importance.GroupTaylorImportance):
-        # Taylor expansion requires gradients for importance estimation
-        # A dummy loss, please replace it with your loss function and data!
-        loss = model(example_inputs).sum() 
-        loss.backward() # before pruner.step()
-    pruner.step()
-    macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
-    # finetune the pruned model here
-    # finetune(model)
-    # ...
+if isinstance(imp, tp.importance.GroupTaylorImportance):
+    # Taylor expansion requires gradients for importance estimation
+    # A dummy loss, please replace it with your loss function and data!
+    loss = model(example_inputs).sum() 
+    loss.backward() # before pruner.step()
+
+pruner.step()
+macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
+# finetune the pruned model here
+# finetune(model)
+# ...
 ```
 #### Global Pruning
 
@@ -204,7 +206,7 @@ for epoch in range(epochs):
 ```
 
 #### Interactive Pruning (Advanced)
-All high-level pruners offer support for interactive pruning. You can utilize the method "pruner.step(interactive=True)" to retrieve all the groups and interactively prune them by calling "group.prune()". This feature is particularly useful when you want to have control over or monitor the pruning process.
+All high-level pruners offer support for interactive pruning. You can utilize the method `pruner.step(interactive=True)` to retrieve all the groups and interactively prune them by calling `group.prune()`. This feature is particularly useful when you want to have control over or monitor the pruning process.
 
 ```python
 for i in range(iterative_steps):
@@ -224,7 +226,7 @@ for i in range(iterative_steps):
 
 #### Group-level Pruning
 
-With DepGraph, it is easy to design some "group-level" criteria to estimate the importance of a whole group rather than a single layer. In Torch-pruning, all pruners work in the group level. Check the following results to see how grouping improves the performance of pruning.
+With DepGraph, it is easy to design some "group-level" criteria to estimate the importance of a whole group rather than a single layer. This feature can be also used to sparsify coupled layers, making all the to-be-pruned parameters consistently unimportant. In Torch-pruning, all pruners work at the group level. Check the following results to see how grouping improves the performance of pruning.
 
 <div align="center">
 <img src="assets/group_sparsity.png" width="80%">
@@ -312,7 +314,7 @@ The following [pruning functions](torch_pruning/pruner/function.py) are availabl
 
 ### 5. Customized Layers
 
-Please refer to [examples/hf_transformers/prune_swin.py](examples/hf_transformers/prune_swin.py), which implements a new pruner for the customized module ``SwinPatchMerging``. A more simple example is available at [tests/test_customized_layer.py](https://github.com/VainF/Torch-Pruning/blob/master/tests/test_customized_layer.py).
+Please refer to [examples/transformers/prune_hf_swin.py](examples/transformers/prune_hf_swin.py), which implements a new pruner for the customized module ``SwinPatchMerging``. A more simple example is available at [tests/test_customized_layer.py](https://github.com/VainF/Torch-Pruning/blob/master/tests/test_customized_layer.py).
 
 ### 6. Benchmarks
 
diff --git a/benchmarks/benchmark_latency.py b/benchmarks/benchmark_latency.py
index 221b9ff..06bd991 100644
--- a/benchmarks/benchmark_latency.py
+++ b/benchmarks/benchmark_latency.py
@@ -23,14 +23,14 @@ def main():
 
     # Before Pruning
     macs, params = tp.utils.count_ops_and_params(model, example_input)
-    latency_mu, latency_std = test_latency(model, example_input)
+    latency_mu, latency_std = estimate_latency(model, example_input)
     # print all with .2f
     print(f"[Iter 0] \tSparsity: 0.00, \tMACs: {macs/1e9:.2f} G, \tParams: {params/1e6:.2f} M, \tLatency: {latency_mu:.2f} ms +- {latency_std:.2f} ms")
 
     for iter in range(iterative_steps):
         pruner.step()
         _macs, _params = tp.utils.count_ops_and_params(model, example_input)
-        latency_mu, latency_std = test_latency(model, example_input)
+        latency_mu, latency_std = estimate_latency(model, example_input)
         current_ch_sparsity = 1 / iterative_steps * (iter + 1)
         print(f"[Iter {iter+1}] \tSparsity: {current_ch_sparsity:.2f}, \tMACs: {_macs/1e9:.2f} G, \tParams: {_params/1e6:.2f} M, \tLatency: {latency_mu:.2f} ms +- {latency_std:.2f} ms")
 
@@ -41,7 +41,7 @@ def main():
         #                _ = model(example_input)
         #print(prof)
 
-def test_latency(model, example_inputs, repetitions=50):
+def estimate_latency(model, example_inputs, repetitions=50):
     import numpy as np
     starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
     timings=np.zeros((repetitions,1))
diff --git a/examples/README.md b/examples/README.md
deleted file mode 100644
index d30a560..0000000
--- a/examples/README.md
+++ /dev/null
@@ -1,8 +0,0 @@
-# Examples
-
-* **hf_transformers**: Prune models from [HuggingFace transformers](https://huggingface.co/docs/transformers/index), including Vision Transformers, Swin Transformers and Bert.
-* **timm_models**: Prune models from [pytorch-image-models (timm)](https://github.com/huggingface/pytorch-image-models).
-* **torchvision_models**: Prune models from [Torchvision](https://pytorch.org/vision/stable/models.html).
-* Yolov7
-* Yolov8
-
diff --git a/examples/transformers/draw_acc_curve.py b/examples/transformers/draw_acc_curve.py
new file mode 100644
index 0000000..f5ec6df
--- /dev/null
+++ b/examples/transformers/draw_acc_curve.py
@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+import os
+# Read such line "Epoch: [250] Total time: ... Test:  Acc@1 77.984 Acc@5 94.114" from a file and record the accuracy.
+plt.style.use('ggplot')
+
+def parse_acc_from_file(file_path):
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+    acc = []
+    for line in lines:
+        if 'Test:  Acc@1' in line:
+            acc.append(float(line.split(' ')[-3]))
+    return acc
+
+log_dict = {
+    'Hessian-uniform': 'output/vit_b_16_pruning_hessian_uniform/train.log',
+    'Taylor-uniform': 'output/vit_b_16_pruning_taylor_uniform/train.log',
+    'Taylor-bottleneck': 'output/vit_b_16_pruning_taylor_bottleneck/train.log',
+    'L1-uniform': 'output/vit_b_16_pruning_l1_uniform/train.log',
+    'L2-uniform': 'output/vit_b_16_pruning_l2_uniform/train.log',
+}
+
+plt.figure(figsize=(8, 4), dpi=200)
+for exp_name, log_path in log_dict.items():
+    acc = parse_acc_from_file(log_path)
+    plt.plot(acc, label=exp_name)
+    print(exp_name, "| Last Epoch:", acc[-1], "| Best Epoch:", max(acc))
+                                                                                                                                                                                                                                                                                                                                                                                                       #plt.plot(acc_random, label='Random-uniform')rplt.xlabel('Epoch')
+plt.ylabel('Accuracy')
+plt.legend()
+plt.legend(loc='lower right')
+
+# change style
+plt.savefig('acc.png')
\ No newline at end of file
diff --git a/examples/transformers/finetune.py b/examples/transformers/finetune.py
index 3ecc18f..1dddd38 100644
--- a/examples/transformers/finetune.py
+++ b/examples/transformers/finetune.py
@@ -8,16 +8,44 @@
 import torch.utils.data
 import torchvision
 import torchvision.transforms
+import torch.nn.functional as F
 import utils
 from sampler import RASampler
 from torch import nn
 from torch.utils.data.dataloader import default_collate
 from torchvision.transforms.functional import InterpolationMode
+
 from transforms import get_mixup_cutmix
 import sys
 
+import timm
+
 import torch_pruning as tp
 
+
+def forward(self, x):
+    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)
+    q, k = self.q_norm(q), self.k_norm(k)
+
+    if self.fused_attn:
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.attn_drop.p,
+        )
+    else:
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+
+    x = x.transpose(1, 2).reshape(B, N, -1)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
 def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args, model_ema=None, scaler=None):
     model.train()
     metric_logger = utils.MetricLogger(delimiter="  ")
@@ -30,6 +58,8 @@ def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, arg
         image, target = image.to(device), target.to(device)
         with torch.cuda.amp.autocast(enabled=scaler is not None):
             output = model(image)
+            if args.is_huggingface:
+                output = output.logits
             loss = criterion(output, target)
 
         optimizer.zero_grad()
@@ -72,6 +102,8 @@ def evaluate(model, criterion, data_loader, device, print_freq=100, log_suffix="
             image = image.to(device, non_blocking=True)
             target = target.to(device, non_blocking=True)
             output = model(image)
+            if args.is_huggingface:
+                output = output.logits
             loss = criterion(output, target)
 
             acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
@@ -140,6 +172,8 @@ def load_data(traindir, valdir, args):
         dataset = torchvision.datasets.ImageFolder(
             traindir,
             presets.ClassificationPresetTrain(
+                mean=(0.485, 0.456, 0.406) if (args.is_huggingface or args.use_imagenet_mean_std) else (0.5, 0.5, 0.5) ,
+                std=(0.229, 0.224, 0.225) if (args.is_huggingface or args.use_imagenet_mean_std) else (0.5, 0.5, 0.5),
                 crop_size=train_crop_size,
                 interpolation=interpolation,
                 auto_augment_policy=auto_augment_policy,
@@ -171,6 +205,8 @@ def load_data(traindir, valdir, args):
 
         else:
             preprocessing = presets.ClassificationPresetEval(
+                mean=(0.485, 0.456, 0.406) if (args.is_huggingface or args.use_imagenet_mean_std) else (0.5, 0.5, 0.5) ,
+                std=(0.229, 0.224, 0.225) if (args.is_huggingface or args.use_imagenet_mean_std) else (0.5, 0.5, 0.5),
                 crop_size=val_crop_size,
                 resize_size=val_resize_size,
                 interpolation=interpolation,
@@ -250,11 +286,21 @@ def collate_fn(batch):
     )
 
     print("Creating model")
-    if args.pruned_model is not None:
-        model = torch.load(args.pruned_model, map_location='cpu') #torchvision.models.get_model(args.model, weights=args.weights, num_classes=num_classes)
-    else:
+    if os.path.isfile(args.model):
+        print(f"Loading pruned model from {args.model}")
+        model = torch.load(args.model, map_location='cpu') #torchvision.models.get_model(args.model, weights=args.weights, num_classes=num_classes)
+    elif args.is_huggingface:
+        print(f"Loading Huggingface model from {args.model}")
         from transformers import ViTForImageClassification
         model = ViTForImageClassification.from_pretrained(args.model)
+    else:
+        print(f"Loading timm model from {args.model}")
+        model = timm.create_model(args.model, pretrained=True)
+    
+    if not args.is_huggingface:
+        for m in model.modules():
+            if isinstance(m, timm.models.vision_transformer.Attention):
+                m.forward = forward.__get__(m, timm.models.vision_transformer.Attention) # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
 
     model.to(device)
     macs, params = tp.utils.count_ops_and_params(model, torch.randn(1, 3, 224, 224).to(device))
@@ -408,8 +454,8 @@ def get_args_parser(add_help=True):
 
     parser = argparse.ArgumentParser(description="PyTorch Classification Training", add_help=add_help)
 
-    parser.add_argument("--model", required=True, type=str, help="model_name")
-    parser.add_argument("--pruned-model", default=None, type=str, help="path to the pruned model")
+    parser.add_argument("--model", default='vit_base_patch16_224', type=str, help="path or name of the (pruned) model")
+    #parser.add_argument("--pruned-model", default=None, type=str, help="path to the pruned model")
     parser.add_argument("--data-path", default="~/Datasets/imagenet", type=str, help="dataset path")
     parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
     parser.add_argument(
@@ -515,7 +561,7 @@ def get_args_parser(add_help=True):
         "--use-deterministic-algorithms", action="store_true", help="Forces the use of deterministic algorithms only."
     )
     parser.add_argument(
-        "--interpolation", default="bilinear", type=str, help="the interpolation method (default: bilinear)"
+        "--interpolation", default="bilinear", type=str, help="the interpolation method (default: bicubic)"
     )
     parser.add_argument(
         "--val-resize-size", default=256, type=int, help="the resize size used for validation (default: 256)"
@@ -535,7 +581,10 @@ def get_args_parser(add_help=True):
     parser.add_argument("--backend", default="PIL", type=str.lower, help="PIL or tensor - case insensitive")
     parser.add_argument("--use-v2", action="store_true", help="Use V2 transforms")
 
+    parser.add_argument("--is_huggingface", action="store_true", help="Use huggingface models")
+
     parser.add_argument("--checkpoint-interval", default=10, type=int, help="checkpoint interval (default: 10)")
+    parser.add_argument("--use_imagenet_mean_std", default=False, action="store_true", help="Use imagenet mean and std")
     return parser
 
 
diff --git a/examples/transformers/prune_bert.py b/examples/transformers/prune_hf_bert.py
similarity index 100%
rename from examples/transformers/prune_bert.py
rename to examples/transformers/prune_hf_bert.py
diff --git a/examples/transformers/prune_swin.py b/examples/transformers/prune_hf_swin.py
similarity index 100%
rename from examples/transformers/prune_swin.py
rename to examples/transformers/prune_hf_swin.py
diff --git a/examples/transformers/prune_hf_vit.py b/examples/transformers/prune_hf_vit.py
index fb43e26..bc79b48 100644
--- a/examples/transformers/prune_hf_vit.py
+++ b/examples/transformers/prune_hf_vit.py
@@ -14,7 +14,7 @@
 
 parser = argparse.ArgumentParser(description='ViT Pruning')
 parser.add_argument('--model_name', default='google/vit-base-patch16-224', type=str, help='model name')
-parser.add_argument('--dataset_root', default='~/Datasets/shared/imagenet/', type=str, help='model name')
+parser.add_argument('--data_path', default='~/Datasets/shared/imagenet/', type=str, help='model name')
 parser.add_argument('--taylor_batchs', default=10, type=int, help='number of batchs for taylor criterion')
 parser.add_argument('--pruning_ratio', default=0.5, type=float, help='prune ratio')
 parser.add_argument('--bottleneck', default=False, action='store_true', help='bottleneck or uniform')
@@ -24,6 +24,7 @@
 
 parser.add_argument('--train_batch_size', default=64, type=int, help='train batch size')
 parser.add_argument('--val_batch_size', default=128, type=int, help='val batch size')
+parser.add_argument('--save_as', default=None, type=str, help='save as')
 args = parser.parse_args()
 
 def prepare_imagenet(imagenet_root, train_batch_size=64, val_batch_size=128, num_workers=4):
@@ -82,7 +83,7 @@ def validate_model(model, val_loader, device):
 else: raise NotImplementedError
 
 if args.pruning_type=='taylor' or args.test_accuracy:
-    train_loader, val_loader = prepare_imagenet(args.dataset_root, train_batch_size=args.train_batch_size, val_batch_size=args.val_batch_size)
+    train_loader, val_loader = prepare_imagenet(args.data_path, train_batch_size=args.train_batch_size, val_batch_size=args.val_batch_size)
 
 # Load the model
 model = ViTForImageClassification.from_pretrained(args.model_name).to(device)
@@ -151,4 +152,10 @@ def validate_model(model, val_loader, device):
 print("Base Params: %.2f M, Pruned Params: %.2f M"%(base_params/1e6, pruned_params/1e6))
 if args.test_accuracy:
     print("Base Loss: %.4f, Pruned Loss: %.4f"%(loss_ori, loss_pruned))
-    print("Base Accuracy: %.4f, Pruned Accuracy: %.4f"%(acc_ori, acc_pruned))
\ No newline at end of file
+    print("Base Accuracy: %.4f, Pruned Accuracy: %.4f"%(acc_ori, acc_pruned))
+
+if args.save_as is not None:
+    print("Saving the pruned model to %s..."%args.save_as)
+    os.makedirs(os.path.dirname(args.save_as), exist_ok=True)
+    model.zero_grad()
+    torch.save(model, args.save_as)
\ No newline at end of file
diff --git a/examples/transformers/prune_timm_vit.py b/examples/transformers/prune_timm_vit.py
index f976c68..f1d8f65 100644
--- a/examples/transformers/prune_timm_vit.py
+++ b/examples/transformers/prune_timm_vit.py
@@ -7,20 +7,23 @@
 from torchvision.datasets import ImageFolder
 import torchvision.transforms as T
 from tqdm import tqdm
+from torchvision.transforms.functional import InterpolationMode
+
+import presets
 
 import argparse
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Timm ViT Pruning')
     parser.add_argument('--model_name', default='vit_base_patch16_224', type=str, help='model name')
-    parser.add_argument('--dataset_root', default='~/Datasets/shared/imagenet/', type=str, help='model name')
+    parser.add_argument('--data_path', default='~/Datasets/shared/imagenet/', type=str, help='model name')
     parser.add_argument('--taylor_batchs', default=10, type=int, help='number of batchs for taylor criterion')
     parser.add_argument('--pruning_ratio', default=0.5, type=float, help='prune ratio')
     parser.add_argument('--bottleneck', default=False, action='store_true', help='bottleneck or uniform')
-    parser.add_argument('--pruning_type', default='l1', type=str, help='pruning type', choices=['random', 'taylor', 'l1'])
+    parser.add_argument('--pruning_type', default='l1', type=str, help='pruning type', choices=['random', 'taylor', 'l2', 'l1', 'hessian'])
     parser.add_argument('--test_accuracy', default=False, action='store_true', help='test accuracy')
     parser.add_argument('--global_pruning', default=False, action='store_true', help='global pruning')
-
+    parser.add_argument('--use_imagenet_mean_std', default=False, action='store_true', help='use imagenet mean and std')
     parser.add_argument('--train_batch_size', default=64, type=int, help='train batch size')
     parser.add_argument('--val_batch_size', default=128, type=int, help='val batch size')
     parser.add_argument('--save_as', default=None, type=str, help='save the pruned model')
@@ -52,32 +55,28 @@ def forward(self, x):
     x = self.proj_drop(x)
     return x
 
-def prepare_imagenet(imagenet_root, train_batch_size=64, val_batch_size=128, num_workers=4):
+def prepare_imagenet(imagenet_root, train_batch_size=64, val_batch_size=128, num_workers=4, use_imagenet_mean_std=False):
     """The imagenet_root should contain train and val folders.
     """
 
     print('Parsing dataset...')
-    train_dst = ImageFolder(os.path.join(imagenet_root, 'train'), transform=T.Compose(
-        [
-            T.Resize(256),
-            T.CenterCrop(224),
-            T.ToTensor(),
-            T.Normalize(
-                mean=[0.485, 0.456, 0.406], 
-                std=[0.229, 0.224, 0.225]
-            ),
-        ])
+    train_dst = ImageFolder(os.path.join(imagenet_root, 'train'), 
+                            transform=presets.ClassificationPresetEval(
+                                mean=[0.485, 0.456, 0.406] if use_imagenet_mean_std else [0.5, 0.5, 0.5],
+                                std=[0.229, 0.224, 0.225] if use_imagenet_mean_std else [0.5, 0.5, 0.5],
+                                crop_size=224,
+                                resize_size=256,
+                                interpolation=InterpolationMode.BILINEAR,
+                            )
     )
-    val_dst = ImageFolder(os.path.join(imagenet_root, 'val'), transform=T.Compose(
-        [
-            T.Resize(256),
-            T.CenterCrop(224),
-            T.ToTensor(),
-            T.Normalize(
-                mean=[0.485, 0.456, 0.406], 
-                std=[0.229, 0.224, 0.225]
-            ),
-        ])
+    val_dst = ImageFolder(os.path.join(imagenet_root, 'val'), 
+                          transform=presets.ClassificationPresetEval(
+                                mean=[0.485, 0.456, 0.406] if use_imagenet_mean_std else [0.5, 0.5, 0.5],
+                                std=[0.229, 0.224, 0.225] if use_imagenet_mean_std else [0.5, 0.5, 0.5],
+                                crop_size=224,
+                                resize_size=256,
+                                interpolation=InterpolationMode.BILINEAR,
+                            )
     )
     train_loader = torch.utils.data.DataLoader(train_dst, batch_size=train_batch_size, shuffle=True, num_workers=num_workers)
     val_loader = torch.utils.data.DataLoader(val_dst, batch_size=val_batch_size, shuffle=False, num_workers=num_workers)
@@ -104,20 +103,23 @@ def main():
     if args.pruning_type == 'random':
         imp = tp.importance.RandomImportance()
     elif args.pruning_type == 'taylor':
-        imp = tp.importance.TaylorImportance()
+        imp = tp.importance.GroupTaylorImportance()
+    elif args.pruning_type == 'l2':
+        imp = tp.importance.GroupNormImportance(p=2)
     elif args.pruning_type == 'l1':
-        imp = tp.importance.MagnitudeImportance(p=1)
+        imp = tp.importance.GroupNormImportance(p=1)
+    elif args.pruning_type == 'hessian':
+        imp = tp.importance.GroupHessianImportance()
     else: raise NotImplementedError
 
-    if args.pruning_type=='taylor' or args.test_accuracy:
-        train_loader, val_loader = prepare_imagenet(args.dataset_root, train_batch_size=args.train_batch_size, val_batch_size=args.val_batch_size)
+    if args.pruning_type in ['taylor', 'hessian'] or args.test_accuracy:
+        train_loader, val_loader = prepare_imagenet(args.data_path, train_batch_size=args.train_batch_size, val_batch_size=args.val_batch_size, use_imagenet_mean_std=args.use_imagenet_mean_std)
 
     # Load the model
     model = timm.create_model(args.model_name, pretrained=True).eval().to(device)
     input_size = [3, 224, 224]
     example_inputs = torch.randn(1, *input_size).to(device)
     base_macs, base_params = tp.utils.count_ops_and_params(model, example_inputs)
-    print(base_macs/1e9, base_params/1e6)
 
     print("Pruning %s..."%args.model_name)
     ch_groups = {}
@@ -142,18 +144,28 @@ def main():
                     ch_sparsity=args.pruning_ratio, # target sparsity
                     ignored_layers=ignored_layers,
                     channel_groups=ch_groups,
+                    round_to=16,
     )
-    if isinstance(imp, tp.importance.TaylorImportance):
+    if isinstance(imp, (tp.importance.GroupTaylorImportance, tp.importance.GroupHessianImportance)):
         model.zero_grad()
-        print("Accumulating gradients for taylor pruning...")
+        if isinstance(imp, tp.importance.GroupHessianImportance):
+            imp.zero_grad()
+        print("Accumulating gradients for pruning...")
         for k, (imgs, lbls) in enumerate(train_loader):
             if k>=args.taylor_batchs: break
             imgs = imgs.to(device)
             lbls = lbls.to(device)
             output = model(imgs)
-            loss = torch.nn.functional.cross_entropy(output, lbls)
-            loss.backward()
-
+            if isinstance(imp, tp.importance.GroupHessianImportance):
+                loss = torch.nn.functional.cross_entropy(output, lbls, reduction='none')
+                for l in loss:
+                    model.zero_grad()
+                    l.backward(retain_graph=True)
+                    imp.accumulate_grad(model)
+            elif isinstance(imp, tp.importance.GroupTaylorImportance):
+                loss = torch.nn.functional.cross_entropy(output, lbls)
+                loss.backward()
+                
     for i, g in enumerate(pruner.step(interactive=True)):
         g.prune()
 
diff --git a/examples/transformers/readme.md b/examples/transformers/readme.md
index b45df87..0be8c15 100644
--- a/examples/transformers/readme.md
+++ b/examples/transformers/readme.md
@@ -1,12 +1,29 @@
-# Examples for HuggingFace Transformers
+# Transformers
 
-Our post-training scripts will be released.
+## Pruning ViT-ImageNet-21K-ft-1K from [Timm](https://github.com/huggingface/pytorch-image-models)
 
-## Pruning ViT from [Timm](https://github.com/huggingface/pytorch-image-models)
+### Data
+Please prepare the ImageNet-1K dataset as follows and modify the data root in the script.
+```
+imagenet/
+  train/
+    n01440764/
+      n01440764_10026.JPEG
+      ...
+    n01773157/
+    n02051845/
+    ...
+  val/
+    n01440764/
+      ILSVRC2012_val_00000293.JPEG
+      ...
+    n01773157/
+    n02051845/
+```
 
-#### Pruning
+### Pruning
 ```bash
-bash scripts/prune_timm_vit_b_16.sh
+bash scripts/prune_timm_vit_b_16_taylor_uniform.sh
 ```
 
 ```
@@ -15,46 +32,98 @@ bash scripts/prune_timm_vit_b_16.sh
 Summary:
 Base MACs: 17.59 G, Pruned MACs: 4.61 G
 Base Params: 86.57 M, Pruned Params: 22.05 M
-Base Loss: 0.7774, Pruned Loss: 7.1335
-Base Accuracy: 0.8108, Pruned Accuracy: 0.0033
-Saving the pruned model to output/pruned/vit_base_patch16_224_pruned.pth...
+Base Loss: 0.6516, Pruned Loss: 7.2412
+Base Accuracy: 0.8521, Pruned Accuracy: 0.0016
+Saving the pruned model to output/pruned/vit_base_patch16_224_pruned_taylor_uniform.pth...
 ```
 
-#### Finetuning
+### Finetuning
 ```bash
-bash scripts/finetune_timm_vit_b_16.sh
+bash scripts/finetune_timm_vit_b_16_taylor_uniform.sh
 ```
-Some results:
+Pruning results for ImageNet-21K-ft-1K (Timm):
 
-| | Vit-B/16 (Timm) | Vit-B/16 (HF) |	ViT_B/32 (Timm) | Group L1 (Uniform) | Group Taylor (Uniform) | Group Taylor (Bottleneck) |
+|  | ViT-B/16 (Timm) |	ViT_B/32 (Timm) | Group L2 (Uniform) | Group Taylor (Uniform) | Group Taylor (Bottleneck) | Group Hessian (Uniform) |
 | :-- | :--: | :--: | :--: | :--: | :--: | :--: |
-| **#Params** | 86.57 M	| 86.56 M	|  	88.22 M | 22.05 M | 22.05 M | 22.8 M |
-| **MACs** | 17.59 G	| 17.59 G	| 4.41 G |  4.61 G	| 4.61 G | 4.23 G |
-| **Acc @ Ep 300** | 81.08	| 75.66	| 72.26 |  79.20	| 79.61 | 79.11 |
-| **Acc @ Ep 50** | -	| -	| - |  69.24	| 71.93 | 71.54 |
+| **#Params** | 86.57 M		|  	88.22 M | 22.05 M | 22.05 M | 24.83 M | 22.05 M |
+| **MACs** | 17.59 G		| 4.41 G |  4.61 G	| 4.61 G | 4.62 G | 4.61 G |
+| **Acc @ Epoch 300** | 85.21	| 80.68  | 78.11 | 80.19 | 80.06 | 80.15   |
+| **Latency (Bs=1, A5000)** | 5.21 ms <br> +- 0.05 ms	|  3.87 ms <br> +- 0.05 ms | 3.99 ms <br> +- 0.10 ms | 3.99 ms <br> +- 0.10 ms  |  3.87 ms <br> +- 0.14 ms  |  3.99 ms <br> +- 0.10 ms    |
+| **Checkpoints** | - | - | [ckpt](https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_l2_uniform.pth) | [ckpt](https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_taylor_uniform.pth) | [ckpt](https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_taylor_bottleneck.pth) | [ckpt](https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_hessian_uniform.pth) |
+
+*Notes:*
+* Uniform - We apply the same pruning ratio to all layers.
+* Bottleneck - We only prune the internal dimensions of Attention & FFN, leading to bottleneck structures.
+* Please adjust the learning rate accordingly if the batch size and number of GPUs are changed. Refer to [this paper](https://arxiv.org/pdf/1706.02677.pdf) for more details about linear LR scaling with large mini-batch.
+
+<div align="center">
+  <img src="https://github.com/VainF/Torch-Pruning/assets/18592211/2537a09f-6d62-4879-8300-6ceec722ebe9" width="100%"></img>
+</div>
 
-* Uniform - The same pruning ratio for all layers.
-* Bottleneck - Only prune the internal dimensions of Attention & FFN.
+### Which pruner should be used for ViT pruning?
 
+In short, `tp.importance.GroupTaylorImportance` + `tp.pruner.MetaPruner` is a good choice for ViT pruning.
+
+* Prune a Vision Transformer (ImageNet-1K) from HF Transformers without fine-tuning.
+<div align="center">
+<img src="https://github.com/VainF/Torch-Pruning/assets/18592211/6f99aa90-259d-41e8-902a-35675a9c9d90" width="45%"></img>
+<img src="https://github.com/VainF/Torch-Pruning/assets/18592211/11473499-d28a-434b-a8d6-1a53c4b3b7c0" width="45%"></img>
+</div>
+
+* Prune a Vision Transformer (ImageNet-21K-ft-1K) from timm without finetuning
 <div align="center">
-<img src="https://github.com/VainF/Torch-Pruning/assets/18592211/24de19ff-60aa-4402-94e3-527670ffb55e" width="80%"></img>
+<img src="https://github.com/VainF/Torch-Pruning/assets/18592211/8726aaf3-129a-4ff6-855e-c73573a8d3e4" width="45%"></img>
+<img src="https://github.com/VainF/Torch-Pruning/assets/18592211/5ff31ebe-3d0e-417b-8020-68afaa19dc65" width="45%"></img>
 </div>
 
+### Latency
+
+* Download our finetuned models
+```bash
+mkdir pretrained
+cd pretrained
+wget https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_taylor_uniform.pth
+wget https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_taylor_bottleneck.pth
+wget https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_l2_uniform.pth
+wget https://github.com/VainF/Torch-Pruning/releases/download/v1.2.5/vit_b_16_pruning_hessian_uniform.pth
+```
 
+* Measure the latency of the pruned models
+```bash
+python estimate_latency.py --model pretrained/vit_b_16_pruning_taylor_uniform.pth
+```
 
-## Pruning Other Transformers
+## Pruning ViT-ImageNet-1K from [HF Transformers](https://huggingface.co/docs/transformers/index)
 
-### ViT from [HF Transformers](https://huggingface.co/docs/transformers/index)
+### Pruning
 ```bash
-python prune_hf_vit.py
+bash scripts/prune_hf_vit_b_16_taylor_uniform.sh  
 ```
 ```
 ...
-Base MACs: 16.848735 G, Pruned MACs: 4.241336 G
-Base Params: 86.567656 M, Pruned Params: 22.050664 M
+----------------------------------------
+Summary:
+Base MACs: 16.85 G, Pruned MACs: 4.24 G
+Base Params: 86.57 M, Pruned Params: 22.05 M
+Base Loss: 0.9717, Pruned Loss: 7.0871
+Base Accuracy: 0.7566, Pruned Accuracy: 0.0015
+Saving the pruned model to output/pruned/hf_vit_base_patch16_224_pruned_taylor_uniform.pth...
 ```
 
-## Swin Transformers from [HF Transformers](https://huggingface.co/docs/transformers/index)
+### Finetuning
+```bash
+bash scripts/finetune_hf_vit_b_16_taylor_uniform.sh
+```
+
+Pruning results for ImageNet-1K (HF Transformers):
+
+| | ViT-B/16 <br> (HF) | ViT-B/16 <br> (Torchvision) |	ViT_B/32 <br> (Torchvision) | Group L1 <br> (Uniform) | Group Taylor <br> (Uniform) | Group Taylor <br> (Bottleneck) |
+| :-- | :--: | :--: | :--: | :--: | :--: | :--: |
+| **#Params** |  86.56 M	| 86.57 M	|  	88.22 M | 22.05 M | 22.05 M | 22.8 M |
+| **MACs** | 17.59 G	| 17.59 G	|  4.41 G |  4.61 G	| 4.61 G | 4.23 G |
+| **Acc @ Ep 300** | 75.66 | 81.068	|  75.91 |  79.20	| 79.61 | 79.11 |
+
+## Pruning Swin Transformers from [HF Transformers](https://huggingface.co/docs/transformers/index)
 ```bash
 python prune_hf_swin.py
 ```
@@ -64,7 +133,7 @@ Base MACs: 4.350805 G, Pruned MACs: 1.438424 G
 Base Params: 28.288354 M, Pruned Params: 9.462802 M
 ```
 
-## Bert from [HF Transformers](https://huggingface.co/docs/transformers/index)
+## Pruning Bert from [HF Transformers](https://huggingface.co/docs/transformers/index)
 ```bash
 python prune_hf_bert.py
 ```
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16.sh b/examples/transformers/scripts/finetune_hf_vit_b_16_l1_uniform.sh
similarity index 74%
rename from examples/transformers/scripts/finetune_timm_vit_b_16.sh
rename to examples/transformers/scripts/finetune_hf_vit_b_16_l1_uniform.sh
index 2670421..cf746e7 100644
--- a/examples/transformers/scripts/finetune_timm_vit_b_16.sh
+++ b/examples/transformers/scripts/finetune_hf_vit_b_16_l1_uniform.sh
@@ -1,6 +1,5 @@
 torchrun --nproc_per_node=8 finetune.py \
-    --pruned-model "output/pruned/vit_base_patch16_224_pruned.pth" \
-    --model google/vit-base-patch16-224 \
+    --model "output/pruned/hf_vit_base_patch16_224_pruned_l1_uniform.pth" \
     --epochs 300 \
     --batch-size 256 \
     --opt adamw \
@@ -18,4 +17,5 @@ torchrun --nproc_per_node=8 finetune.py \
     --ra-sampler \
     --cutmix-alpha 1.0 \
     --data-path "~/Datasets/shared/imagenet/" \
-    --output-dir output/vit_b_16_taylor_pruning
\ No newline at end of file
+    --output-dir output/hf_vit_b_16_pruning_l1_uniform \
+    --is_huggingface \
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_hf_vit_b_16_taylor_uniform.sh b/examples/transformers/scripts/finetune_hf_vit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..bc192e2
--- /dev/null
+++ b/examples/transformers/scripts/finetune_hf_vit_b_16_taylor_uniform.sh
@@ -0,0 +1,21 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/hf_vit_base_patch16_224_pruned_taylor_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/hf_vit_b_16_pruning_taylor_uniform \
+    --is_huggingface \
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_deit_b_16_taylor_uniform.sh b/examples/transformers/scripts/finetune_timm_deit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..bf5dd16
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_deit_b_16_taylor_uniform.sh
@@ -0,0 +1,22 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/deit_base_patch16_224_pruned_taylor_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --random-erase 0.25 \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/deit_b_16_pruning_taylor_uniform \
+    --use_imagenet_mean_std \
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16_hessian_uniform.sh b/examples/transformers/scripts/finetune_timm_vit_b_16_hessian_uniform.sh
new file mode 100644
index 0000000..779f14b
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_vit_b_16_hessian_uniform.sh
@@ -0,0 +1,20 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/vit_base_patch16_224_pruned_hessian_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/vit_b_16_pruning_hessian_uniform
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16_l1_uniform.sh b/examples/transformers/scripts/finetune_timm_vit_b_16_l1_uniform.sh
new file mode 100644
index 0000000..f66a0e8
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_vit_b_16_l1_uniform.sh
@@ -0,0 +1,20 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/vit_base_patch16_224_pruned_l1_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/vit_b_16_pruning_l1_uniform
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16_l2_uniform.sh b/examples/transformers/scripts/finetune_timm_vit_b_16_l2_uniform.sh
new file mode 100644
index 0000000..1e764fa
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_vit_b_16_l2_uniform.sh
@@ -0,0 +1,20 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/vit_base_patch16_224_pruned_l2_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/vit_b_16_pruning_l2_uniform
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_bottleneck.sh b/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_bottleneck.sh
new file mode 100644
index 0000000..0cfb755
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_bottleneck.sh
@@ -0,0 +1,20 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/vit_base_patch16_224_pruned_taylor_bottleneck.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/vit_b_16_pruning_taylor_bottleneck
\ No newline at end of file
diff --git a/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_uniform.sh b/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..888085f
--- /dev/null
+++ b/examples/transformers/scripts/finetune_timm_vit_b_16_taylor_uniform.sh
@@ -0,0 +1,21 @@
+torchrun --nproc_per_node=8 finetune.py \
+    --model "output/pruned/vit_base_patch16_224_pruned_taylor_uniform.pth" \
+    --epochs 300 \
+    --batch-size 256 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --random-erase 0.25 \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --output-dir output/vit_b_16_pruning_taylor_uniform_v2
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_hf_vit_b_16_l1_uniform.sh b/examples/transformers/scripts/prune_hf_vit_b_16_l1_uniform.sh
new file mode 100644
index 0000000..8a3e3bd
--- /dev/null
+++ b/examples/transformers/scripts/prune_hf_vit_b_16_l1_uniform.sh
@@ -0,0 +1,10 @@
+python prune_hf_vit.py \
+    --model_name google/vit-base-patch16-224 \
+    --pruning_type l1 \
+    --pruning_ratio 0.5 \
+    --taylor_batchs 10 \
+    --data_path ~/Datasets/shared/imagenet \
+    --test_accuracy \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/hf_vit_base_patch16_224_pruned_taylor_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_hf_vit_b_16_taylor_uniform.sh b/examples/transformers/scripts/prune_hf_vit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..3d60650
--- /dev/null
+++ b/examples/transformers/scripts/prune_hf_vit_b_16_taylor_uniform.sh
@@ -0,0 +1,10 @@
+python prune_hf_vit.py \
+    --model_name google/vit-base-patch16-224 \
+    --pruning_type taylor \
+    --pruning_ratio 0.5 \
+    --taylor_batchs 10 \
+    --data_path ~/Datasets/shared/imagenet \
+    --test_accuracy \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/hf_vit_base_patch16_224_pruned_taylor_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_deit_b_16_taylor_uniform.sh b/examples/transformers/scripts/prune_timm_deit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..d58e94d
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_deit_b_16_taylor_uniform.sh
@@ -0,0 +1,10 @@
+python prune_timm_vit.py \
+    --model_name deit_base_distilled_patch16_224 \
+    --pruning_type taylor \
+    --pruning_ratio 0.54 \
+    --taylor_batchs 50 \
+    --data_path ~/Datasets/shared/imagenet \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/deit_base_patch16_224_pruned_taylor_uniform.pth \
+    --use_imagenet_mean_std \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16_hessian_uniform.sh b/examples/transformers/scripts/prune_timm_vit_b_16_hessian_uniform.sh
new file mode 100644
index 0000000..02d9b29
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_hessian_uniform.sh
@@ -0,0 +1,10 @@
+python prune_timm_vit.py \
+    --model_name vit_base_patch16_224 \
+    --pruning_type hessian \
+    --pruning_ratio 0.5 \
+    --taylor_batchs 10 \
+    --test_accuracy \
+    --data_path ~/Datasets/shared/imagenet \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/vit_base_patch16_224_pruned_hessian_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16.sh b/examples/transformers/scripts/prune_timm_vit_b_16_l1_uniform.sh
similarity index 57%
rename from examples/transformers/scripts/prune_timm_vit_b_16.sh
rename to examples/transformers/scripts/prune_timm_vit_b_16_l1_uniform.sh
index 545681b..7cdb17b 100644
--- a/examples/transformers/scripts/prune_timm_vit_b_16.sh
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_l1_uniform.sh
@@ -1,10 +1,10 @@
 python prune_timm_vit.py \
     --model_name vit_base_patch16_224 \
-    --pruning_type taylor \
+    --pruning_type l1 \
     --pruning_ratio 0.5 \
     --taylor_batchs 10 \
-    --dataset_root ~/Datasets/shared/imagenet \
+    --data_path ~/Datasets/shared/imagenet \
     --test_accuracy \
     --train_batch_size 64 \
     --val_batch_size 64 \
-    --save_as output/pruned/vit_base_patch16_224_pruned.pth \
\ No newline at end of file
+    --save_as output/pruned/vit_base_patch16_224_pruned_l1_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16_l2_uniform.sh b/examples/transformers/scripts/prune_timm_vit_b_16_l2_uniform.sh
new file mode 100644
index 0000000..cdd51e4
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_l2_uniform.sh
@@ -0,0 +1,10 @@
+python prune_timm_vit.py \
+    --model_name vit_base_patch16_224 \
+    --pruning_type l2 \
+    --pruning_ratio 0.5 \
+    --taylor_batchs 10 \
+    --data_path ~/Datasets/shared/imagenet \
+    --test_accuracy \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/vit_base_patch16_224_pruned_l2_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16_taylor_bottleneck.sh b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_bottleneck.sh
new file mode 100644
index 0000000..149238c
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_bottleneck.sh
@@ -0,0 +1,10 @@
+python prune_timm_vit.py \
+    --model_name vit_base_patch16_224 \
+    --pruning_type taylor \
+    --pruning_ratio 0.73 \
+    --taylor_batchs 10 \
+    --data_path ~/Datasets/shared/imagenet \
+    --bottleneck \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/vit_base_patch16_224_pruned_taylor_bottleneck.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform.sh b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform.sh
new file mode 100644
index 0000000..2763a36
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform.sh
@@ -0,0 +1,9 @@
+python prune_timm_vit.py \
+    --model_name vit_base_patch16_224 \
+    --pruning_type taylor \
+    --pruning_ratio 0.54 \
+    --taylor_batchs 50 \
+    --data_path ~/Datasets/shared/imagenet \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/vit_base_patch16_224_pruned_taylor_uniform.pth \
\ No newline at end of file
diff --git a/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform_global.sh b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform_global.sh
new file mode 100644
index 0000000..9173b20
--- /dev/null
+++ b/examples/transformers/scripts/prune_timm_vit_b_16_taylor_uniform_global.sh
@@ -0,0 +1,10 @@
+python prune_timm_vit.py \
+    --model_name vit_base_patch16_224 \
+    --pruning_type taylor \
+    --pruning_ratio 0.6 \
+    --taylor_batchs 10 \
+    --data_path ~/Datasets/shared/imagenet \
+    --train_batch_size 64 \
+    --val_batch_size 64 \
+    --save_as output/pruned/vit_base_patch16_224_pruned_taylor_uniform.pth \
+    --global_pruning
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 54ec9da..c114d29 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="torch-pruning",
-    version="v1.2.4",
+    version="v1.2.5",
     author="Gongfan Fang",
     author_email="gongfan@u.nus.edu",
     description="Towards Any Structural Pruning",


Test patch:
diff --git a/examples/transformers/scripts/test_pretrained_hf_vit_b_16.sh b/examples/transformers/scripts/test_pretrained_hf_vit_b_16.sh
new file mode 100644
index 0000000..353b0d0
--- /dev/null
+++ b/examples/transformers/scripts/test_pretrained_hf_vit_b_16.sh
@@ -0,0 +1,22 @@
+python finetune.py \
+    --model "google/vit-base-patch16-224" \
+    --epochs 300 \
+    --batch-size 32 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --test-only \
+    --interpolation bilinear \
+    --is_huggingface \
\ No newline at end of file
diff --git a/examples/transformers/scripts/test_pretrained_timm_deit_b_16.sh b/examples/transformers/scripts/test_pretrained_timm_deit_b_16.sh
new file mode 100644
index 0000000..9d6e356
--- /dev/null
+++ b/examples/transformers/scripts/test_pretrained_timm_deit_b_16.sh
@@ -0,0 +1,21 @@
+python finetune.py \
+    --model "deit_base_distilled_patch16_224" \
+    --epochs 300 \
+    --batch-size 32 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --test-only \
+    --use_imagenet_mean_std \
diff --git a/examples/transformers/scripts/test_pretrained_timm_vit_b_16.sh b/examples/transformers/scripts/test_pretrained_timm_vit_b_16.sh
new file mode 100644
index 0000000..08b021e
--- /dev/null
+++ b/examples/transformers/scripts/test_pretrained_timm_vit_b_16.sh
@@ -0,0 +1,20 @@
+python finetune.py \
+    --model "vit_base_patch16_224" \
+    --epochs 300 \
+    --batch-size 32 \
+    --opt adamw \
+    --lr 0.00015 \
+    --wd 0.3 \
+    --lr-scheduler cosineannealinglr \
+    --lr-warmup-method linear \
+    --lr-warmup-epochs 0 \
+    --lr-warmup-decay 0.033 \
+    --amp \
+    --label-smoothing 0.11 \
+    --mixup-alpha 0.2 \
+    --auto-augment ra \
+    --clip-grad-norm 1 \
+    --ra-sampler \
+    --cutmix-alpha 1.0 \
+    --data-path "~/Datasets/shared/imagenet/" \
+    --test-only \
diff --git a/examples/transformers/test_latency.py b/examples/transformers/test_latency.py
new file mode 100644
index 0000000..59e5da0
--- /dev/null
+++ b/examples/transformers/test_latency.py
@@ -0,0 +1,81 @@
+import sys, os
+import time
+sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+import torch_pruning as tp
+import torch
+import timm
+import torch.nn.functional as F
+
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', type=str, default='vit_base_patch16_224', help='model name or path')
+parser.add_argument('--batch_size', type=int, default=1, help='batch size')
+args = parser.parse_args()
+
+def forward(self, x):
+    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)
+    q, k = self.q_norm(q), self.k_norm(k)
+
+    if self.fused_attn:
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.attn_drop.p,
+        )
+    else:
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+
+    x = x.transpose(1, 2).reshape(B, N, -1)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if os.path.isfile(args.model):
+        loaded_pth = torch.load(args.model, map_location='cpu')
+        if isinstance(loaded_pth, dict):
+            model = loaded_pth['model'].to(device)
+        else:
+            model = loaded_pth.to(device)
+    else:
+        model = timm.create_model(args.model, pretrained=True).to(device)
+        
+    for m in model.modules():
+        if isinstance(m, timm.models.vision_transformer.Attention):
+            m.forward = forward.__get__(m, timm.models.vision_transformer.Attention) # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
+
+    example_input = torch.rand(args.batch_size, 3, 224, 224).to(device)
+    macs, params = tp.utils.count_ops_and_params(model, example_input)
+    latency_mu, latency_std = estimate_latency(model, example_input)
+    print(f"MACs: {macs/1e9:.2f} G, \tParams: {params/1e6:.2f} M, \tLatency: {latency_mu:.2f} ms +- {latency_std:.2f} ms")
+
+def estimate_latency(model, example_inputs, repetitions=300):
+    import numpy as np
+    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    timings=np.zeros((repetitions,1))
+
+    for _ in range(50):
+        _ = model(example_inputs)
+
+    with torch.no_grad():
+        for rep in range(repetitions):
+            starter.record()
+            _ = model(example_inputs)
+            ender.record()
+            # WAIT FOR GPU SYNC
+            torch.cuda.synchronize()
+            curr_time = starter.elapsed_time(ender)
+            timings[rep] = curr_time
+
+    mean_syn = np.sum(timings) / repetitions
+    std_syn = np.std(timings)
+    return mean_syn, std_syn
+
+if __name__=='__main__':
+    main()
\ No newline at end of file