Package: torch-pruning
Towards Any Structural Pruning[Documentation & Tutorials] [FAQ]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, includingLarge Language Models (LLMs),Segment Anything Model (SAM),Diffusion Models,Yolov7,yolov8,Vision Transformers,Swin Transformers,BERT, FasterRCNN, SSD, ResNe(X)t, ConvNext, DenseNet, ConvNext, RegNet, DeepLab, etc. Different fromtorch.nn.utils.prunethat zeroizes parameters through masking, Torch-Pruning deploys an algorithm calledDepGraphto remove parameters physically.Examples: Pruning off-the-shelf models from Timm, Huggingface Transformers, Torchvision, Yolo, etc.Benchmark: Reproduce the our results in the DepGraph paper.For more technical details, please refer to our CVPR'23 paper:DepGraph: Towards Any Structural PruningGongfan Fang,Xinyin Ma,Mingli Song,Michael Bi Mi,Xinchao WangLearning and Vision Lab, National University of SingaporeUpdate:2023.12.19 :rocket:IntroducingDeepCache: Accelerating Diffusion Models for Free2023.12.19 :rocket:SlimSAM: 0.1% Data Makes Segment Anything Slim2023.09.06Pruning & Finetuning Examples for Vision Transformers, Swin Transformers, Bert2023.07.19 Support LLaMA, LLaMA-2, Vicuna, Baichuan, Bloom inLLM-Pruner2023.05.20LLM-Pruner: On the Structural Pruning of Large Language Models[arXiv]2023.05.19Structural Pruning for Diffusion Models[arXiv]2023.04.15Pruning and Post-training for YOLOv7 / YOLOv8Features:High-level Pruners:MetaPruner,MagnitudePruner,BNScalePruner,GroupNormPruner,GrowingRegPruner, RandomPruner, etc. A paper list is available on ourwiki page.Dependency Graph for automatic structural pruningLow-level pruning functionsSupported Importance Criteria: L-p Norm, Taylor, Random, BNScaling, etc.Supported modules: Linear, (Transposed) Conv, Normalization, PReLU, Embedding, MultiheadAttention, nn.Parameters,customized modulesand nested/composed modules.Supported operators: split, concatenation, skip connection, flatten, reshape, view, all element-wise ops, etc.Benchmarks,TutorialsandExamplesContact Us:Please do not hesitate to open anissueif you encounter any problems with the library or the paper.Or Join our Discord or WeChat group for a chat:Discord:linkWeChat Group (Group size exceeded 400):QR CodeTable of ContentsInstallationQuickstartHow It WorksA Minimal Example of DepGraphHigh-level PrunersGlobal PruningSparse TrainingInteractive PruningSoft PruningGroup-level PruningModify Module Attributes or Forward FunctionSave & LoadLow-level Pruning FunctionsCustomized LayersBenchmarksOur Results on {ResNet-56 / CIFAR-10 / 2.00x}LatencySeries of WorksCitationInstallationTorch-Pruning is compatible with both PyTorch 1.x and 2.x versions. However, it is highly recommended to use PyTorch 2.0.pipinstalltorch-pruningorgitclonehttps://github.com/VainF/Torch-Pruning.gitQuickstartHere we provide a quick start for Torch-Pruning. More explained details can be found inTutoralsHow It WorksIn structural pruning, a "Group" is defined as the minimal removable unit within deep networks. Most groups are composed of multiple layers that are interdependent and need to be pruned together in order to maintain the integrity of the resulting structures. However, deep networks often have complex dependencies among their layers, making structural pruning a challenging task. This work addresses this challenge by introducing an automated mechanism called "DepGraph." DepGraph allows for seamless parameter grouping and facilitates pruning in various types of deep networks.A Minimal Example of DepGraphPlease ensure that your model is set up to enable AutoGrad withouttorch.no_grador.requires_grad=False.importtorchfromtorchvision.modelsimportresnet18importtorch_pruningastpmodel=resnet18(pretrained=True).eval()# 1. Build dependency graph for resnet18DG=tp.DependencyGraph().build_dependency(model,example_inputs=torch.randn(1,3,224,224))# 2. Group coupled layers for model.conv1group=DG.get_pruning_group(model.conv1,tp.prune_conv_out_channels,idxs=[2,6,9])# 3. Prune grouped layers altogetherifDG.check_pruning_group(group):# avoid full pruning, i.e., channels=0.group.prune()# 4. Save & Loadmodel.zero_grad()# clear gradientstorch.save(model,'model.pth')# We can not use .state_dict as the model structure is changed.model=torch.load('model.pth')# load the pruned modelThe above example demonstrates the basic pruning pipeline with DepGraph. The target layerresnet.conv1is coupled with multiple layers, necessitating their simultaneous removal during structural pruning. To observe the cascading effect of pruning operations, we can print the groups and observe how one pruning operation can "trigger" others. In the subsequent outputs, "A => B" indicates that pruning operation "A" triggers pruning operation "B." The group[0] refers to the pruning root in DG.get_pruning_group. For more details about grouping, please refer toWiki - DepGraph & Group.--------------------------------
          Pruning Group
--------------------------------
[0] prune_out_channels on conv1 (Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on conv1 (Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)), idxs=[2, 6, 9] (Pruning Root)
[1] prune_out_channels on conv1 (Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), idxs=[2, 6, 9]
[2] prune_out_channels on bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_20(ReluBackward0), idxs=[2, 6, 9]
[3] prune_out_channels on _ElementWiseOp_20(ReluBackward0) => prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward0), idxs=[2, 6, 9]
[4] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward0) => prune_out_channels on _ElementWiseOp_18(AddBackward0), idxs=[2, 6, 9]
[5] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward0) => prune_in_channels on layer1.0.conv1 (Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
[6] prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on layer1.0.bn2 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), idxs=[2, 6, 9]
[7] prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on _ElementWiseOp_17(ReluBackward0), idxs=[2, 6, 9]
[8] prune_out_channels on _ElementWiseOp_17(ReluBackward0) => prune_out_channels on _ElementWiseOp_16(AddBackward0), idxs=[2, 6, 9]
[9] prune_out_channels on _ElementWiseOp_17(ReluBackward0) => prune_in_channels on layer1.1.conv1 (Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
[10] prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on layer1.1.bn2 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), idxs=[2, 6, 9]
[11] prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on _ElementWiseOp_15(ReluBackward0), idxs=[2, 6, 9]
[12] prune_out_channels on _ElementWiseOp_15(ReluBackward0) => prune_in_channels on layer2.0.downsample.0 (Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)), idxs=[2, 6, 9]
[13] prune_out_channels on _ElementWiseOp_15(ReluBackward0) => prune_in_channels on layer2.0.conv1 (Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
[14] prune_out_channels on layer1.1.bn2 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.1.conv2 (Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
[15] prune_out_channels on layer1.0.bn2 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.0.conv2 (Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
--------------------------------How to scan all groups (Advanced):We can useDG.get_all_groups(ignored_layers, root_module_types)to scan and prune all groups sequentially. Each group will begin with a layer that matches one type in theroot_module_typesparameter. Note thatDG.get_all_groupsis only responsible for grouping and does not have any knowledge or understanding of which parameters should be pruned. Therefore, it is necessary to specify the pruning idxs usinggroup.prune(idxs=idxs). This feature is useful when you want to implement your own pruning algorithms.forgroupinDG.get_all_groups(ignored_layers=[model.conv1],root_module_types=[nn.Conv2d,nn.Linear]):# handle groups in sequential orderidxs=[2,4,6]# your pruning indicesgroup.prune(idxs=idxs)print(group)High-level PrunersWith DepGraph, we developed several high-level pruners in this repository to facilitate effortless pruning. By specifying the desired channel pruning ratio, a pruner will scan all prunable groups, estimate the importance, prune the entire model, and fine-tune it using your own training code. For detailed information on this process, please refer tothis tutorial, which shows how to implement aslimmingpruner from scratch. Additionally, a more practical example is available inbenchmarks/main.py.importtorchfromtorchvision.modelsimportresnet18importtorch_pruningastpmodel=resnet18(pretrained=True)example_inputs=torch.randn(1,3,224,224)# 1. Importance criterionimp=tp.importance.GroupTaylorImportance()# or GroupNormImportance(p=2), GroupHessianImportance(), etc.# 2. Initialize a pruner with the model and the importance criterionignored_layers=[]forminmodel.modules():ifisinstance(m,torch.nn.Linear)andm.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,pruning_ratio=0.5,# remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}# pruning_ratio_dict = {model.conv1: 0.2, model.layer2: 0.8}, # customized pruning ratios for layers or blocksignored_layers=ignored_layers,)# 3. Prune & finetune the modelbase_macs,base_nparams=tp.utils.count_ops_and_params(model,example_inputs)ifisinstance(imp,tp.importance.GroupTaylorImportance):# Taylor expansion requires gradients for importance estimationloss=model(example_inputs).sum()# A dummy loss, please replace this line with your loss function and data!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 PruningWith the option of global pruning (global_pruning=True), adaptive sparsity will be allocated to different layers based on their global rank of importance. While this strategy can offer performance advantages, it also carries the potential of overly pruning specific layers, resulting in a substantial decline in overall performance.If you're not very familiar with pruning, it's recommended to begin withglobal_pruning=False.Sparse TrainingSome pruners likeBNScalePrunerandGroupNormPrunersupport sparse training. This can be easily achieved by insertingpruner.update_regularizer()before training, andpruner.regularize(model)betweenloss.backward()andoptimizer.step(). The pruner will accumulate the regularization gradients to.grad.forepochinrange(epochs):model.train()pruner.update_regularizer()# <== initialize regularizerfori,(data,target)inenumerate(train_loader):data,target=data.to(device),target.to(device)optimizer.zero_grad()out=model(data)loss=F.cross_entropy(out,target)loss.backward()# after loss.backward()pruner.regularize(model)# <== for sparse trainingoptimizer.step()# before optimizer.step()Interactive PruningAll high-level pruners offer support for interactive pruning. You can utilize the methodpruner.step(interactive=True)to retrieve all the groups and interactively prune them by callinggroup.prune(). This feature is particularly useful when you want to have control over or monitor the pruning process.foriinrange(iterative_steps):forgroupinpruner.step(interactive=True):# Warning: groups must be handled sequentially. Do not keep them as a list.print(group)# do whatever you like with the groupdep,idxs=group[0]# get the idxstarget_module=dep.target.module# get the root modulepruning_fn=dep.handler# get the pruning functiongroup.prune()# group.prune(idxs=[0, 2, 6]) # It is even possible to change the pruning behaviour with the idxs parametermacs,nparams=tp.utils.count_ops_and_params(model,example_inputs)# finetune your model here# finetune(model)# ...Soft PruningIt is easy to implement Soft Pruning leveraginginteractive=True, which zeros out parameters without removing them. An example can be found intests/test_soft_pruning.pyGroup-level PruningWith 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 sparse. In Torch-pruning, all pruners work at the group level. Check the following results to see how grouping improves the performance of pruning.Pruning a ResNet50 pre-trained on ImageNet-1K without fine-tuning.Pruning a Vision Transformer pre-trained on ImageNet-1K without fine-tuning.Modify module attributes or forward functionIn some implementation, model forward might rely on some static attributes. For example inconvformer_s18of timm, we have:classScale(nn.Module):"""Scale vector by element multiplications."""def__init__(self,dim,init_value=1.0,trainable=True,use_nchw=True):super().__init__()self.shape=(dim,1,1)ifuse_nchwelse(dim,)# static shape, which should be updated after pruningself.scale=nn.Parameter(init_value*torch.ones(dim),requires_grad=trainable)defforward(self,x):returnx*self.scale.view(self.shape)# => x * self.scale.view(-1, 1, 1), this works for pruningwhere theforwardfunction relies onself.shapeduring forwarding. But, the trueself.shapechanged after pruning, which should be manually updated.Save and LoadMethod 1:The following script saves the whole model object (structure+weights) as a 'model.pth'.model.zero_grad()# Remove gradientstorch.save(model,'model.pth')# without .state_dictmodel=torch.load('model.pth')# load the pruned modelMethod 2 (Experimental Features):Re-create pruned models from unpruned ones usingtp.state_dictandtp.load_state_dict.# save the pruned state_dict, which includes both pruned parameters and modified attributesstate_dict=tp.state_dict(pruned_model)# the pruned model, e.g., a resnet-18-halftorch.save(state_dict,'pruned.pth')# create a new model, e.g. resnet18new_model=resnet18().eval()# load the pruned state_dict into the unpruned model.loaded_state_dict=torch.load('pruned.pth',map_location='cpu')tp.load_state_dict(new_model,state_dict=loaded_state_dict)Refer totests/test_serialization.pyfor an ViT example. In this example, we will prune the model and modify some attributes likemodel.hidden_dims.Low-level Pruning FunctionsAlthough it is possible to manually prune your model using low-level functions, this approach can be cumbersome and time-consuming due to the need for meticulous management of dependencies. Therefore, we strongly recommend utilizing the high-level pruners mentioned earlier to streamline and simplify the pruning process. These pruners provide a more convenient and efficient way to perform pruning on your models. To manually prune themodel.conv1of a ResNet-18, the pruning pipeline should look like this:tp.prune_conv_out_channels(model.conv1,idxs=[2,6,9])# fix the broken dependencies manuallytp.prune_batchnorm_out_channels(model.bn1,idxs=[2,6,9])tp.prune_conv_in_channels(model.layer2[0].conv1,idxs=[2,6,9])...The followingpruning functionsare available:'prune_conv_out_channels','prune_conv_in_channels','prune_depthwise_conv_out_channels','prune_depthwise_conv_in_channels','prune_batchnorm_out_channels','prune_batchnorm_in_channels','prune_linear_out_channels','prune_linear_in_channels','prune_prelu_out_channels','prune_prelu_in_channels','prune_layernorm_out_channels','prune_layernorm_in_channels','prune_embedding_out_channels','prune_embedding_in_channels','prune_parameter_out_channels','prune_parameter_in_channels','prune_multihead_attention_out_channels','prune_multihead_attention_in_channels','prune_groupnorm_out_channels','prune_groupnorm_in_channels','prune_instancenorm_out_channels','prune_instancenorm_in_channels',Customized LayersPlease refer toexamples/transformers/prune_hf_swin.py, which implements a new pruner for the customized moduleSwinPatchMerging. A more simple example is available attests/test_customized_layer.py.BenchmarksOur results on {ResNet-56 / CIFAR-10 / 2.00x}MethodBase (%)Pruned (%)$\Delta$ Acc (%)Speed UpNIPS[1]---0.031.76xGeometric[2]93.5993.26-0.331.70xPolar[3]93.8093.83+0.031.88xCP[4]92.8091.80-1.002.00xAMC[5]92.8091.90-0.902.00xHRank[6]93.2692.17-0.092.00xSFP[7]93.5993.36+0.232.11xResRep[8]93.7193.71+0.002.12xOurs-L193.5392.93-0.602.12xOurs-BN93.5393.29-0.242.12xOurs-Group93.5393.77+0.382.13xLatencyLatency test on ResNet-50, Batch Size=64.[Iter 0]        Pruning ratio: 0.00,         MACs: 4.12 G,   Params: 25.56 M,        Latency: 45.22 ms +- 0.03 ms
[Iter 1]        Pruning ratio: 0.05,         MACs: 3.68 G,   Params: 22.97 M,        Latency: 46.53 ms +- 0.06 ms
[Iter 2]        Pruning ratio: 0.10,         MACs: 3.31 G,   Params: 20.63 M,        Latency: 43.85 ms +- 0.08 ms
[Iter 3]        Pruning ratio: 0.15,         MACs: 2.97 G,   Params: 18.36 M,        Latency: 41.22 ms +- 0.10 ms
[Iter 4]        Pruning ratio: 0.20,         MACs: 2.63 G,   Params: 16.27 M,        Latency: 39.28 ms +- 0.20 ms
[Iter 5]        Pruning ratio: 0.25,         MACs: 2.35 G,   Params: 14.39 M,        Latency: 34.60 ms +- 0.19 ms
[Iter 6]        Pruning ratio: 0.30,         MACs: 2.02 G,   Params: 12.46 M,        Latency: 33.38 ms +- 0.27 ms
[Iter 7]        Pruning ratio: 0.35,         MACs: 1.74 G,   Params: 10.75 M,        Latency: 31.46 ms +- 0.20 ms
[Iter 8]        Pruning ratio: 0.40,         MACs: 1.50 G,   Params: 9.14 M,         Latency: 29.04 ms +- 0.19 ms
[Iter 9]        Pruning ratio: 0.45,         MACs: 1.26 G,   Params: 7.68 M,         Latency: 27.47 ms +- 0.28 ms
[Iter 10]       Pruning ratio: 0.50,         MACs: 1.07 G,   Params: 6.41 M,         Latency: 20.68 ms +- 0.13 ms
[Iter 11]       Pruning ratio: 0.55,         MACs: 0.85 G,   Params: 5.14 M,         Latency: 20.48 ms +- 0.21 ms
[Iter 12]       Pruning ratio: 0.60,         MACs: 0.67 G,   Params: 4.07 M,         Latency: 18.12 ms +- 0.15 ms
[Iter 13]       Pruning ratio: 0.65,         MACs: 0.53 G,   Params: 3.10 M,         Latency: 15.19 ms +- 0.01 ms
[Iter 14]       Pruning ratio: 0.70,         MACs: 0.39 G,   Params: 2.28 M,         Latency: 13.47 ms +- 0.01 ms
[Iter 15]       Pruning ratio: 0.75,         MACs: 0.29 G,   Params: 1.61 M,         Latency: 10.07 ms +- 0.01 ms
[Iter 16]       Pruning ratio: 0.80,         MACs: 0.18 G,   Params: 1.01 M,         Latency: 8.96 ms +- 0.02 ms
[Iter 17]       Pruning ratio: 0.85,         MACs: 0.10 G,   Params: 0.57 M,         Latency: 7.03 ms +- 0.04 ms
[Iter 18]       Pruning ratio: 0.90,         MACs: 0.05 G,   Params: 0.25 M,         Latency: 5.81 ms +- 0.03 ms
[Iter 19]       Pruning ratio: 0.95,         MACs: 0.01 G,   Params: 0.06 M,         Latency: 5.70 ms +- 0.03 ms
[Iter 20]       Pruning ratio: 1.00,         MACs: 0.01 G,   Params: 0.06 M,         Latency: 5.71 ms +- 0.03 msPlease refer tobenchmarksfor more details.Series of WorksDepGraph: Towards Any Structural Pruning[Project][Paper]Gongfan Fang, Xinyin Ma, Mingli Song, Michael Bi Mi, Xinchao WangCVPR 2023LLM-Pruner: On the Structural Pruning of Large Language Models[Project][arXiv]Xinyin Ma, Gongfan Fang, Xinchao WangNeurIPS 2023Structural Pruning for Diffusion Models[Project][arxiv]Gongfan Fang, Xinyin Ma, Xinchao WangNeurIPS 2023DeepCache: Accelerating Diffusion Models for Free[Project][Arxiv]Xinyin Ma, Gongfan Fang, and Xinchao WangPreprint 20230.1% Data Makes Segment Anything Slim[Project][Arxiv]Zigeng Chen, Gongfan Fang, Xinyin Ma, Xinchao WangPreprint 2023Citation@inproceedings{fang2023depgraph,
  title={Depgraph: Towards any structural pruning},
  author={Fang, Gongfan and Ma, Xinyin and Song, Mingli and Mi, Michael Bi and Wang, Xinchao},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={16091--16101},
  year={2023}
}