2. Quick Start
下文我们提供了有关Torch-Pruning的简单教程。
0. How it Works
在本项目的结构化剪枝中，组被定义成深度神经网络中可移除的最小单位。每一个组包含了多个相互关联的层，这些层需要同时被剪枝以保持剪枝后网络结构的完整性。然而，深度神经网络通常具有非常复杂的层间依赖关系，这为结构化剪枝带来了非常巨大的挑战。为了解决这一问题，本研究给出了一种自动化解决方案叫做DepGraph，能够将参数高效分组，并为多种常用的神经网络结构的剪枝提供支持，从而简化了剪枝过程。
1. What is Dependency
Naive pruning
为了说明什么是依赖性，我们以ResNet-18模型为例进行结构化剪枝。以下样例代码尝试从conv1层中移除索引为0，1的通道：
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True).eval()
tp.prune_conv_out_channels(model.conv1, idxs=[0,1]) # remove channel 0 and channel 1
output = model(torch.randn(1,3,224,224)) # test
然而，在这种剪枝后的模型，会出现以下无效的网络结构：
ResNet(
(conv1): Conv2d(3, 62, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
...
在前向传播的过程中会出现尺寸不匹配的错误：
RuntimeError: running_mean should contain 62 elements not 64
在这个样例中，可以很轻松的发现只要把bn1层也应该被剪枝掉。因此，可以认为model.conv1和model.bn1存在依赖关系。此外，后续卷积层的输入通道数（即model.layer1[0]）也必须进行相应的调整。
An improved version
实际上，上述案例的依赖关系比我们现在所观察到的更加复杂。我们继续改进代码，如果处理了BN和Conv之后会发生什么。
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True).eval()
tp.prune_conv_out_channels(model.conv1, idxs=[0,1])
tp.prune_batchnorm_out_channels(model.bn1, idxs=[0,1])
tp.prune_batchnorm_in_channels(model.layer1[0].conv1, idxs=[0,1])
output = model(torch.randn(1,3,224,224))
尽管代码经过这些优化，但是由于残差连接的特殊结构，运行代码仍然会发生尺寸不匹配的错误。
File "/home_local/xxx/miniconda3/lib/python3.9/site-packages/torchvision/models/resnet.py", line 102, in forward
out += identity
RuntimeError: The size of tensor a (64) must match the size of tensor b (62) at non-singleton dimension 1
因此，实际情况中结构化剪枝实际特别复杂。Torch-Pruning提供了一种通用且简单的方式，自动化地处理这些依赖关系，从而显著简化剪枝过程。
2. A Minimal Example
使用DepGraph进行剪枝:
import torch
from torchvision.models import resnet18
import torch_pruning as tp
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. Specify the to-be-pruned channels. Here we prune those channels indexed by [2, 6, 9].
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
# 3. prune all grouped layers that are coupled with model.conv1 (included).
if DG.check_pruning_group(group): # avoid full pruning, i.e., channels=0.
group.prune()
# 4. Save & Load
model.zero_grad() # We don't want to store gradient information
torch.save(model, 'model.pth') # without .state_dict
model = torch.load('model.pth') # load the model object
上述示例代码展示了利用DepGraph的简单的剪枝流程。要剪枝的目标层conv1与多个层之间存在耦合关系，在结构化剪枝中，需要同时剪枝掉这些层。我们打印出剪枝组（Group），观察一次剪枝操作如何“触发”其他剪枝操作。A=>B表示剪枝操作A触发了剪枝操作B。group[0]是DG.get_pruning_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):
我们可以使用DG.get_all_groups(ignored_layers, root_module_types)按顺序扫描所有剪枝组。每个组的起始层都必须与root_module_types参数中指定的类型匹配。需要注意的是，DG.get_all_groups仅仅用来分组，并不直接进行剪枝操作。具体的操作通过group.prune(idxs=idxs)指定所需要剪枝的通道。
for group in DG.get_all_groups(ignored_layers=[model.conv1], root_module_types=[nn.Conv2d, nn.Linear]):
# handle groups in sequential order
idxs = [2,4,6] # your pruning indices
group.prune(idxs=idxs)
print(group)
3. High-level Pruners
Pruning with High-level Pruners
利用DependencyGraph，我们开发了许多高级剪枝工具，可实现轻松剪枝。通过指定的通道稀疏度，可以对整个模型进行剪枝，然后使用自己的训练代码进行微调。
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True)
# Importance criteria
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.TaylorImportance()
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!
iterative_steps = 5 # progressive pruning
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
iterative_steps=iterative_steps,
ch_sparsity=0.5, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
ignored_layers=ignored_layers,
)
base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
for i in range(iterative_steps):
if isinstance(imp, tp.importance.TaylorImportance):
# Taylor expansion requires gradients for importance estimation
loss = model(example_inputs).sum() # a dummy loss for TaylorImportance
loss.backward() # before pruner.step()
pruner.step()
macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
# finetune your model here
# finetune(model)
# ...
Sparse Training
某些剪枝工具（如BNScalePruner和GroupNormPruner）在剪枝前需要稀疏训练。这可以通过训练脚本中添加一行代码实现：pruner.regularize(model)。剪枝工具会自动更新可训练参数的梯度。
for epoch in range(epochs):
model.train()
for i, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
out = model(data)
loss = F.cross_entropy(out, target)
loss.backward()
pruner.regularize(model) # <== for sparse learning
optimizer.step()
Interactive Pruning (Advanced)
所有高级剪枝工具都支持交互式剪枝。通过pruner.step(interactive=True)获取所有剪枝组（Group），并调用group.prune()对其逐步剪枝。此功能适合需要控制或监控剪枝过程的情况。
for i in range(iterative_steps):
for group in pruner.step(interactive=True): # Warning: groups must be handled sequentially. Do not keep them as a list.
print(group)
# do whatever you like with the group
dep, idxs = group[0] # get the idxs
target_module = dep.target.module # get the root module
pruning_fn = dep.handler # get the pruning function
# Don't forget to prune the group
group.prune()
# group.prune(idxs=[0, 2, 6]) # It is even possible to change the pruning behaviour with the idxs parameter
macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
# finetune your model here
# finetune(model)
# ...
Group-level Pruning
通过DepGraph，可以轻松设计基于“组级别”的准则来评估整个剪枝组的重要性，而非单个层。Torch-pruning中的所有剪枝工具均以组为单位运作。
4. Save & Load
以下脚本将整个模型对象（结构+权重）保存为model.pth
model.zero_grad() # We don't want to store gradient information
torch.save(model, 'model.pth') # without .state_dict
model = torch.load('model.pth') # load the pruned model
实验性功能：从未剪枝模型中重新创建剪枝后的模型
# save the pruned state_dict, which includes both pruned parameters and modified attributes
state_dict = tp.state_dict(pruned_model) # the pruned model, e.g., a resnet-18-half
torch.save(state_dict, 'pruned.pth')
# create a new model, e.g. resnet18
new_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)
print(new_model) # This will be a pruned model.
5. Low-level Pruning Functions
虽然可以使用低级函数手动剪枝模型，但这种方法非常繁琐，因为需要仔细管理相关依赖关系。因此，建议使用前述的高级剪枝工具以简化剪枝过程。
tp.prune_conv_out_channels( model.conv1, idxs=[2,6,9] )
# fix the broken dependencies manually
tp.prune_batchnorm_out_channels( model.bn1, idxs=[2,6,9] )
tp.prune_conv_in_channels( model.layer2[0].conv1, idxs=[2,6,9] )
...
剪枝函数如下：
'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',
6. Customized Layers
参考Torch-Pruning/tests/test_customized_layer.py at master · VainF/Torch-Pruning
3. DepGraph & Group
What is DepGraph
Dependency Graph（DepGraph）是Torch-pruning的核心功能之一，它为分组依赖层提供了自动化机制。DepGraph包含两个关键概念：
tp.dependency.Dependency：层与层之间的依赖关系tp.dependency.DependencyGraph:用于建模依赖关系的依赖图tp.dependency.Group:表示剪枝的最小操作单位，里面是层与层的依赖列表
以下教程以ResNet-18为例，展示如何使用DepGraph。
Pruning Pipeline with DepGraph
在Torch-Pruning中，DepGraph会找到与要剪枝的层相依赖的层，并归为一组。首先通过输入和模型构建DepGraph，追踪模型来得到模型的计算图。然后，选择要剪枝的层通过idxs参数指定需要剪枝的通道。
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True).eval()
# 1. Build the DepGraph with an example input.
DG = tp.DependencyGraph().build_dependency(model, example_inputs=torch.randn(1,3,224,224))
# 2. Specify the to-be-pruned channels. Here we prune those channels indexed by [2, 6, 9].
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
# 3. Prune all coupled layers in this group
group.prune()
output = model(torch.randn(1,3,224,224)) # Test
Manipulating a Group
Get a group
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
# Choose 1. a root layer, 2. a pruning function, 3. a index list
Print a group
print(group.details()) # use print(group) if you are not interested in the full idxs list.
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on conv1 (Conv2d(3, 61, 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, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), idxs=[2, 6, 9]
[2] prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_20(ReluBackward1), idxs=[2, 6, 9]
[3] prune_out_channels on _ElementWiseOp_20(ReluBackward1) => prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward), idxs=[2, 6, 9]
[4] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_out_channels on _ElementWiseOp_18(AddBackward0), idxs=[2, 6, 9]
[5] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_in_channels on layer1.0.conv1 (Conv2d(61, 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(61, 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(ReluBackward1), idxs=[2, 6, 9]
[8] prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_out_channels on _ElementWiseOp_16(AddBackward0), idxs=[2, 6, 9]
[9] prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_in_channels on layer1.1.conv1 (Conv2d(61, 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(61, 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(ReluBackward1), idxs=[2, 6, 9]
[12] prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.downsample.0 (Conv2d(61, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)), idxs=[2, 6, 9]
[13] prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.conv1 (Conv2d(61, 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(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.1.conv2 (Conv2d(64, 61, 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(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.0.conv2 (Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), idxs=[2, 6, 9]
--------------------------------
Prune a group
Pruning with the pre-defined idxs:
print(group[0].idxs)
group.prune()
[2, 6, 9]
Pruning with new idxs:
new_idxs = [1,2,3,4]
group.prune(new_idxs)
Iterate a group
for i, (dep, idxs) in enumerate(group):
print("Dep: ", dep, " Idxs:", idxs)
Dep: prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) Idxs: [2, 6, 9]
Dep: prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_20(ReluBackward1) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_20(ReluBackward1) => prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_out_channels on _ElementWiseOp_18(AddBackward0) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_in_channels on layer1.0.conv1 (Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on layer1.0.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on _ElementWiseOp_17(ReluBackward1) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_out_channels on _ElementWiseOp_16(AddBackward0) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_in_channels on layer1.1.conv1 (Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on layer1.1.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on _ElementWiseOp_15(ReluBackward1) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.downsample.0 (Conv2d(61, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.conv1 (Conv2d(61, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on layer1.1.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.1.conv2 (Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) Idxs: [2, 6, 9]
Dep: prune_out_channels on layer1.0.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.0.conv2 (Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) Idxs: [2, 6, 9]
Get the layer & pruning function
for i, (dep, idxs) in enumerate(group):
layer = dep.layer
pruning_fn = dep.pruning_fn
print(layer, pruning_fn)
Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7f339752f7c0>>
_ElementWiseOp_20(ReluBackward1) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
_ElementWiseOp_19(MaxPool2DWithIndicesBackward) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
_ElementWiseOp_18(AddBackward0) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7f339752f7c0>>
_ElementWiseOp_17(ReluBackward1) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
_ElementWiseOp_16(AddBackward0) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7f339752f7c0>>
_ElementWiseOp_15(ReluBackward1) <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7f3396ec5820>>
Conv2d(61, 128, kernel_size=(1, 1), stride=(2, 2), bias=False) <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
Conv2d(61, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f339752f730>>
如果你对"Trigger"的过程感兴趣的话，见以下示例：
for i, (dep, idxs) in enumerate(group):
trigger = dep.trigger
handler = dep.handler
source_layer = dep.source.module
target_layer = dep.target.module
print("For Dep: ", dep)
print(" > Trigger: ", trigger) // Trigger：触发剪枝的函数
print(" > Handler: ", handler) // Handler: 当前层中实际执行的剪枝函数
print(" > Source Layer: ", source_layer) // 触发的源层
print(" > Target Layer: ", target_layer) // 目标层
print("")
这里的 handler 和 target_layer 分别与 dep.pruning_fn 和 dep.layer 相同。
For Dep: prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False))
> Trigger: <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Handler: <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
> Target Layer: Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
For Dep: prune_out_channels on conv1 (Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))
> Trigger: <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Handler: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Source Layer: Conv2d(3, 61, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
> Target Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
For Dep: prune_out_channels on bn1 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_20(ReluBackward1)
> Trigger: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
> Target Layer: _ElementWiseOp_20(ReluBackward1)
For Dep: prune_out_channels on _ElementWiseOp_20(ReluBackward1) => prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward)
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: _ElementWiseOp_20(ReluBackward1)
> Target Layer: _ElementWiseOp_19(MaxPool2DWithIndicesBackward)
For Dep: prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_out_channels on _ElementWiseOp_18(AddBackward0)
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: _ElementWiseOp_19(MaxPool2DWithIndicesBackward)
> Target Layer: _ElementWiseOp_18(AddBackward0)
For Dep: prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_in_channels on layer1.0.conv1 (Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: _ElementWiseOp_19(MaxPool2DWithIndicesBackward)
> Target Layer: Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
For Dep: prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on layer1.0.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Source Layer: _ElementWiseOp_18(AddBackward0)
> Target Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
For Dep: prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on _ElementWiseOp_17(ReluBackward1)
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: _ElementWiseOp_18(AddBackward0)
> Target Layer: _ElementWiseOp_17(ReluBackward1)
For Dep: prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_out_channels on _ElementWiseOp_16(AddBackward0)
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: _ElementWiseOp_17(ReluBackward1)
> Target Layer: _ElementWiseOp_16(AddBackward0)
For Dep: prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_in_channels on layer1.1.conv1 (Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: _ElementWiseOp_17(ReluBackward1)
> Target Layer: Conv2d(61, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
For Dep: prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on layer1.1.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Source Layer: _ElementWiseOp_16(AddBackward0)
> Target Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
For Dep: prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on _ElementWiseOp_15(ReluBackward1)
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Source Layer: _ElementWiseOp_16(AddBackward0)
> Target Layer: _ElementWiseOp_15(ReluBackward1)
For Dep: prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.downsample.0 (Conv2d(61, 128, kernel_size=(1, 1), stride=(2, 2), bias=False))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: _ElementWiseOp_15(ReluBackward1)
> Target Layer: Conv2d(61, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
For Dep: prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.conv1 (Conv2d(61, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False))
> Trigger: <bound method DummyPruner.prune_out_channels of <torch_pruning.ops.ElementWisePruner object at 0x7fe4a1d49eb0>>
> Handler: <bound method ConvPruner.prune_in_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: _ElementWiseOp_15(ReluBackward1)
> Target Layer: Conv2d(61, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
For Dep: prune_out_channels on layer1.1.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.1.conv2 (Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))
> Trigger: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Handler: <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
> Target Layer: Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
For Dep: prune_out_channels on layer1.0.bn2 (BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.0.conv2 (Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))
> Trigger: <bound method BatchnormPruner.prune_out_channels of <torch_pruning.pruner.function.BatchnormPruner object at 0x7fe4a23cce50>>
> Handler: <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7fe4a23ccdc0>>
> Source Layer: BatchNorm2d(61, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
> Target Layer: Conv2d(64, 61, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
Scan all groups
for g in DG.get_all_groups():
print(g)
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on fc (Linear(in_features=512, out_features=1000, bias=True)) => prune_out_channels on fc (Linear(in_features=512, out_features=1000, bias=True)), #idxs=1000
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer4.0.downsample.0 (Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer4.0.downsample.0 (Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=512
[1] prune_out_channels on layer4.0.downsample.0 (Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer4.0.downsample.1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=512
[2] prune_out_channels on layer4.0.downsample.1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_6(AddBackward0), #idxs=512
[3] prune_out_channels on _ElementWiseOp_6(AddBackward0) => prune_out_channels on layer4.0.bn2 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=512
[4] prune_out_channels on _ElementWiseOp_6(AddBackward0) => prune_out_channels on _ElementWiseOp_5(ReluBackward1), #idxs=512
[5] prune_out_channels on _ElementWiseOp_5(ReluBackward1) => prune_out_channels on _ElementWiseOp_4(AddBackward0), #idxs=512
[6] prune_out_channels on _ElementWiseOp_5(ReluBackward1) => prune_in_channels on layer4.1.conv1 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
[7] prune_out_channels on _ElementWiseOp_4(AddBackward0) => prune_out_channels on layer4.1.bn2 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=512
[8] prune_out_channels on _ElementWiseOp_4(AddBackward0) => prune_out_channels on _ElementWiseOp_3(ReluBackward1), #idxs=512
[9] prune_out_channels on _ElementWiseOp_3(ReluBackward1) => prune_out_channels on _ElementWiseOp_2(MeanBackward1), #idxs=512
[10] prune_out_channels on _ElementWiseOp_2(MeanBackward1) => prune_out_channels on _Reshape_0(), #idxs=512
[11] prune_out_channels on _Reshape_0() => prune_in_channels on fc (Linear(in_features=512, out_features=1000, bias=True)), #idxs=512
[12] prune_in_channels on fc (Linear(in_features=512, out_features=1000, bias=True)) => prune_out_channels on _ElementWiseOp_1(TBackward), #idxs=512
[13] prune_out_channels on layer4.1.bn2 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer4.1.conv2 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
[14] prune_out_channels on layer4.0.bn2 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer4.0.conv2 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer3.0.downsample.0 (Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer3.0.downsample.0 (Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=256
[1] prune_out_channels on layer3.0.downsample.0 (Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer3.0.downsample.1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=256
[2] prune_out_channels on layer3.0.downsample.1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_10(AddBackward0), #idxs=256
[3] prune_out_channels on _ElementWiseOp_10(AddBackward0) => prune_out_channels on layer3.0.bn2 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=256
[4] prune_out_channels on _ElementWiseOp_10(AddBackward0) => prune_out_channels on _ElementWiseOp_9(ReluBackward1), #idxs=256
[5] prune_out_channels on _ElementWiseOp_9(ReluBackward1) => prune_out_channels on _ElementWiseOp_8(AddBackward0), #idxs=256
[6] prune_out_channels on _ElementWiseOp_9(ReluBackward1) => prune_in_channels on layer3.1.conv1 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
[7] prune_out_channels on _ElementWiseOp_8(AddBackward0) => prune_out_channels on layer3.1.bn2 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=256
[8] prune_out_channels on _ElementWiseOp_8(AddBackward0) => prune_out_channels on _ElementWiseOp_7(ReluBackward1), #idxs=256
[9] prune_out_channels on _ElementWiseOp_7(ReluBackward1) => prune_in_channels on layer4.0.downsample.0 (Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=256
[10] prune_out_channels on _ElementWiseOp_7(ReluBackward1) => prune_in_channels on layer4.0.conv1 (Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=256
[11] prune_out_channels on layer3.1.bn2 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer3.1.conv2 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
[12] prune_out_channels on layer3.0.bn2 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer3.0.conv2 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer2.0.downsample.0 (Conv2d(58, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer2.0.downsample.0 (Conv2d(58, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=128
[1] prune_out_channels on layer2.0.downsample.0 (Conv2d(58, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)) => prune_out_channels on layer2.0.downsample.1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=128
[2] prune_out_channels on layer2.0.downsample.1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_14(AddBackward0), #idxs=128
[3] prune_out_channels on _ElementWiseOp_14(AddBackward0) => prune_out_channels on layer2.0.bn2 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=128
[4] prune_out_channels on _ElementWiseOp_14(AddBackward0) => prune_out_channels on _ElementWiseOp_13(ReluBackward1), #idxs=128
[5] prune_out_channels on _ElementWiseOp_13(ReluBackward1) => prune_out_channels on _ElementWiseOp_12(AddBackward0), #idxs=128
[6] prune_out_channels on _ElementWiseOp_13(ReluBackward1) => prune_in_channels on layer2.1.conv1 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
[7] prune_out_channels on _ElementWiseOp_12(AddBackward0) => prune_out_channels on layer2.1.bn2 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=128
[8] prune_out_channels on _ElementWiseOp_12(AddBackward0) => prune_out_channels on _ElementWiseOp_11(ReluBackward1), #idxs=128
[9] prune_out_channels on _ElementWiseOp_11(ReluBackward1) => prune_in_channels on layer3.0.downsample.0 (Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=128
[10] prune_out_channels on _ElementWiseOp_11(ReluBackward1) => prune_in_channels on layer3.0.conv1 (Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=128
[11] prune_out_channels on layer2.1.bn2 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer2.1.conv2 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
[12] prune_out_channels on layer2.0.bn2 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer2.0.conv2 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on conv1 (Conv2d(3, 58, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on conv1 (Conv2d(3, 58, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)), #idxs=58
[1] prune_out_channels on conv1 (Conv2d(3, 58, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)) => prune_out_channels on bn1 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=58
[2] prune_out_channels on bn1 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_20(ReluBackward1), #idxs=58
[3] prune_out_channels on _ElementWiseOp_20(ReluBackward1) => prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward), #idxs=58
[4] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_out_channels on _ElementWiseOp_18(AddBackward0), #idxs=58
[5] prune_out_channels on _ElementWiseOp_19(MaxPool2DWithIndicesBackward) => prune_in_channels on layer1.0.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=58
[6] prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on layer1.0.bn2 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=58
[7] prune_out_channels on _ElementWiseOp_18(AddBackward0) => prune_out_channels on _ElementWiseOp_17(ReluBackward1), #idxs=58
[8] prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_out_channels on _ElementWiseOp_16(AddBackward0), #idxs=58
[9] prune_out_channels on _ElementWiseOp_17(ReluBackward1) => prune_in_channels on layer1.1.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=58
[10] prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on layer1.1.bn2 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=58
[11] prune_out_channels on _ElementWiseOp_16(AddBackward0) => prune_out_channels on _ElementWiseOp_15(ReluBackward1), #idxs=58
[12] prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.downsample.0 (Conv2d(58, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)), #idxs=58
[13] prune_out_channels on _ElementWiseOp_15(ReluBackward1) => prune_in_channels on layer2.0.conv1 (Conv2d(58, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=58
[14] prune_out_channels on layer1.1.bn2 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.1.conv2 (Conv2d(64, 58, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=58
[15] prune_out_channels on layer1.0.bn2 (BatchNorm2d(58, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on layer1.0.conv2 (Conv2d(64, 58, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=58
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer1.0.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer1.0.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=64
[1] prune_out_channels on layer1.0.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer1.0.bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=64
[2] prune_out_channels on layer1.0.bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_21(ReluBackward1), #idxs=64
[3] prune_out_channels on _ElementWiseOp_21(ReluBackward1) => prune_in_channels on layer1.0.conv2 (Conv2d(64, 58, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=64
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer1.1.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer1.1.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=64
[1] prune_out_channels on layer1.1.conv1 (Conv2d(58, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer1.1.bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=64
[2] prune_out_channels on layer1.1.bn1 (BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_22(ReluBackward1), #idxs=64
[3] prune_out_channels on _ElementWiseOp_22(ReluBackward1) => prune_in_channels on layer1.1.conv2 (Conv2d(64, 58, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=64
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer2.0.conv1 (Conv2d(58, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer2.0.conv1 (Conv2d(58, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=128
[1] prune_out_channels on layer2.0.conv1 (Conv2d(58, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer2.0.bn1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=128
[2] prune_out_channels on layer2.0.bn1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_23(ReluBackward1), #idxs=128
[3] prune_out_channels on _ElementWiseOp_23(ReluBackward1) => prune_in_channels on layer2.0.conv2 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer2.1.conv1 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer2.1.conv1 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
[1] prune_out_channels on layer2.1.conv1 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer2.1.bn1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=128
[2] prune_out_channels on layer2.1.bn1 (BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_24(ReluBackward1), #idxs=128
[3] prune_out_channels on _ElementWiseOp_24(ReluBackward1) => prune_in_channels on layer2.1.conv2 (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=128
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer3.0.conv1 (Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer3.0.conv1 (Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=256
[1] prune_out_channels on layer3.0.conv1 (Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer3.0.bn1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=256
[2] prune_out_channels on layer3.0.bn1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_25(ReluBackward1), #idxs=256
[3] prune_out_channels on _ElementWiseOp_25(ReluBackward1) => prune_in_channels on layer3.0.conv2 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer3.1.conv1 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer3.1.conv1 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
[1] prune_out_channels on layer3.1.conv1 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer3.1.bn1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=256
[2] prune_out_channels on layer3.1.bn1 (BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_26(ReluBackward1), #idxs=256
[3] prune_out_channels on _ElementWiseOp_26(ReluBackward1) => prune_in_channels on layer3.1.conv2 (Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=256
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer4.0.conv1 (Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer4.0.conv1 (Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)), #idxs=512
[1] prune_out_channels on layer4.0.conv1 (Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)) => prune_out_channels on layer4.0.bn1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=512
[2] prune_out_channels on layer4.0.bn1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_27(ReluBackward1), #idxs=512
[3] prune_out_channels on _ElementWiseOp_27(ReluBackward1) => prune_in_channels on layer4.0.conv2 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
--------------------------------
--------------------------------
Pruning Group
--------------------------------
[0] prune_out_channels on layer4.1.conv1 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer4.1.conv1 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
[1] prune_out_channels on layer4.1.conv1 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)) => prune_out_channels on layer4.1.bn1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)), #idxs=512
[2] prune_out_channels on layer4.1.bn1 (BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)) => prune_out_channels on _ElementWiseOp_28(ReluBackward1), #idxs=512
[3] prune_out_channels on _ElementWiseOp_28(ReluBackward1) => prune_in_channels on layer4.1.conv2 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)), #idxs=512
--------------------------------
A more complicated example: record your pruning process
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18()
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.MagnitudeImportance(p=2)
ignored_layers = []
# DO NOT prune the final classifier!
for m in model.modules():
if isinstance(m, torch.nn.Linear) and m.out_features == 1000:
ignored_layers.append(m)
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
iterative_steps=1,
ch_sparsity=0.2, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
ignored_layers=ignored_layers,
)
records = []
for g in pruner.step(interactive=True):
dep, idxs = g[0]
layer = dep.layer
pruning_fn = dep.pruning_fn
records.append((layer, idxs, pruning_fn))
g.prune()
for rec in records:
print(rec)
print("")
(Conv2d(204, 409, kernel_size=(1, 1), stride=(2, 2), bias=False), [311, 139, 434, 102, 226, 402, 347, 24, 368, 451, 485, 408, 65, 34, 421, 152, 47, 116, 260, 2, 506, 367, 280, 66, 81, 171, 291, 53, 101, 439, 36, 244, 56, 365, 381, 189, 69, 204, 186, 112, 170, 23, 384, 471, 214, 255, 195, 13, 394, 206, 108, 190, 224, 220, 122, 431, 443, 144, 349, 177, 207, 390, 289, 374, 176, 92, 351, 76, 217, 382, 464, 249, 283, 133, 320, 131, 111, 209, 124, 160, 355, 30, 126, 10, 369, 110, 405, 242, 305, 62, 366, 343, 502, 119, 293, 297, 151, 254, 136, 129, 1, 328, 96], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(102, 204, kernel_size=(1, 1), stride=(2, 2), bias=False), [203, 5, 8, 16, 80, 54, 92, 95, 2, 253, 193, 99, 60, 224, 142, 55, 90, 11, 200, 45, 215, 43, 39, 59, 197, 115, 194, 107, 20, 13, 119, 72, 139, 84, 191, 210, 204, 83, 65, 62, 185, 214, 175, 235, 228, 216, 88, 138, 182, 22, 86, 159], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(51, 102, kernel_size=(1, 1), stride=(2, 2), bias=False), [7, 87, 105, 99, 1, 43, 97, 9, 26, 58, 76, 3, 77, 28, 5, 64, 95, 45, 74, 13, 47, 21, 98, 22, 90, 48], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(3, 51, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False), [26, 31, 63, 28, 52, 59, 3, 18, 4, 51, 2, 30, 11], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(51, 51, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), [36, 45, 4, 58, 44, 47, 8, 30, 46, 16, 28, 27, 9], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(51, 51, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), [47, 37, 56, 48, 8, 9, 52, 31, 36, 12, 50, 38, 51], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(51, 102, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False), [44, 124, 88, 83, 60, 63, 14, 113, 54, 90, 30, 114, 110, 127, 111, 40, 87, 11, 81, 98, 121, 100, 61, 36, 45, 58], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(102, 102, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), [118, 61, 82, 1, 51, 86, 37, 70, 113, 58, 66, 125, 79, 55, 122, 72, 22, 48, 7, 50, 41, 99, 107, 21, 6, 23], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(102, 204, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False), [24, 13, 100, 89, 56, 216, 51, 29, 117, 102, 2, 239, 181, 34, 120, 104, 165, 182, 46, 196, 161, 93, 175, 195, 77, 63, 42, 22, 251, 194, 53, 249, 6, 105, 92, 81, 132, 201, 20, 75, 154, 128, 235, 207, 152, 112, 136, 59, 241, 149, 103, 236], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(204, 204, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), [28, 226, 103, 108, 9, 30, 89, 149, 49, 105, 199, 10, 195, 4, 247, 20, 84, 133, 245, 181, 88, 232, 218, 229, 233, 209, 210, 65, 109, 46, 47, 94, 179, 136, 137, 48, 193, 164, 157, 189, 212, 62, 188, 113, 59, 139, 161, 194, 31, 184, 141, 168], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(204, 409, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False), [219, 326, 509, 359, 463, 505, 159, 496, 311, 106, 488, 145, 60, 323, 81, 482, 284, 168, 243, 207, 401, 445, 441, 85, 510, 174, 328, 487, 200, 411, 449, 161, 375, 169, 259, 27, 389, 215, 454, 15, 125, 124, 452, 410, 148, 327, 291, 262, 472, 293, 356, 166, 308, 172, 355, 374, 462, 273, 113, 32, 38, 392, 191, 483, 442, 435, 477, 340, 307, 51, 209, 139, 37, 456, 163, 319, 347, 430, 105, 313, 155, 208, 30, 149, 310, 47, 8, 183, 271, 24, 63, 244, 203, 225, 469, 100, 267, 398, 65, 147, 151, 280, 199], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
(Conv2d(409, 409, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), [32, 360, 223, 423, 242, 17, 146, 177, 503, 433, 430, 141, 40, 281, 214, 163, 197, 263, 65, 15, 471, 257, 392, 48, 493, 347, 268, 329, 339, 51, 72, 20, 499, 106, 510, 27, 86, 221, 293, 427, 482, 422, 224, 96, 57, 442, 179, 2, 298, 272, 374, 125, 139, 41, 408, 209, 54, 248, 303, 366, 113, 93, 175, 312, 128, 473, 73, 99, 142, 264, 42, 75, 33, 235, 330, 437, 396, 195, 464, 414, 461, 498, 205, 8, 36, 487, 384, 276, 496, 237, 147, 181, 418, 368, 393, 297, 67, 306, 456, 165, 296, 231, 495], <bound method ConvPruner.prune_out_channels of <torch_pruning.pruner.function.ConvPruner object at 0x7f33583a9d30>>)
4. High-level Pruners
What is High-level Pruner
在Torch-Pruning中，每种算法都被实现为一个高层次的剪枝器（pruner)，负责执行剪枝过程。剪枝器主要包含三个功能：稀疏训练（可选）、重要性评估和参数移除。为了支持剪枝过程，Torch-Pruning提供了两个核心功能：
- tp.importance()：用于衡量权重重要性的标准。
- tp.pruner(): 用于实际剪枝参数的剪枝器。
Pruning Pipeline with Pruner
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True)
# Importance criteria
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.TaylorImportance()
# Ignore some layers, e.g., the output layer
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!
# Initialize a pruner
iterative_steps = 5 # progressive pruning
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
iterative_steps=iterative_steps,
pruning_ratio=0.5, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
ignored_layers=ignored_layers,
)
# prune the model, iteratively if necessary.
base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
for i in range(iterative_steps):
# Taylor expansion requires gradients for importance estimation
if isinstance(imp, tp.importance.TaylorImportance):
# 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 your model here
# finetune(model)
# ...
Dive into tp.pruner.MetaPruner
Definition
tp.pruner.MetaPruner提供了基础的剪枝功能，并包含以下参数。
class MetaPruner:
def __init__(
self,
# Basic
model: nn.Module, # a simple pytorch model
example_inputs: torch.Tensor, # a dummy input for graph tracing. Should be on the same
importance: typing.Callable, # tp.importance.Importance for group importance estimation
global_pruning: bool = False, # https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#global-pruning.
pruning_ratio: float = 0.5, # channel/dim sparsity
pruning_ratio_dict: typing.Dict[nn.Module, float] = None, # layer-specific sparsity, will cover pruning_ratio if specified
max_pruning_ratio: float = 1.0, # maximum sparsity. useful if over-pruning happens.
iterative_steps: int = 1, # for iterative pruning
iterative_sparsity_scheduler: typing.Callable = linear_scheduler, # scheduler for iterative pruning.
ignored_layers: typing.List[nn.Module] = None, # ignored layers
round_to: int = None, # round channels to a multiple of round_to
# Advanced
channel_groups: typing.Dict[nn.Module, int] = dict(), # channel groups for layers like group convs & group norms
customized_pruners: typing.Dict[typing.Any, function.BasePruningFunc] = None, # pruners for customized layers. E.g., {nn.Linear: my_linear_pruner}
unwrapped_parameters: typing.Dict[nn.Parameter, int] = None, # unwrapped nn.Parameters & pruning_dims. For example, {ViT.pos_emb: 0}
root_module_types: typing.List = [ops.TORCH_CONV, ops.TORCH_LINEAR, ops.TORCH_LSTM], # root module for each group
forward_fn: typing.Callable = None, # a function to execute model.forward
output_transform: typing.Callable = None, # a function to transform network outputs
):
Basic Arguments
model, example_inputs & importance
剪枝器至少需要三个参数才能进行剪枝
global_pruning
https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#global-pruning
pruning_ratio & pruning_ratio_dict
tp.pruner.MetaPruner为我们提供了稀疏性（剪枝比例）的控制，可以通过全局参数pruning_ratio和针对特定层的pruning_ratio_dict来实现。后者接受字典格式，例如{model.block1: 0.2}。如果在pruning_ratio_dict中未显式定义某层的稀疏性，则全局参数pruning_ratio将应用于所有层。
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18()
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.MagnitudeImportance(p=2)
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
imp,
pruning_ratio = 0.5,
pruning_ratio_dict = {model.layer2: 0.2}
)
pruner.step()
print(model)
这里我们为第二个残差块自定义了剪枝比例，这将导致模型结构会发生如下变化：
ResNet{64, 128, 256, 512} => ResNet{32, 102, 128, 256}
ResNet(
(conv1): Conv2d(3, 32, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer2): Sequential(
(0): BasicBlock(
(conv1): Conv2d(32, 102, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(102, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(102, 102, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(102, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(32, 102, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(102, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(102, 102, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(102, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(102, 102, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(102, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer3): Sequential(
(0): BasicBlock(
(conv1): Conv2d(102, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(102, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer4): Sequential(
(0): BasicBlock(
(conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=256, out_features=500, bias=True)
)
max_pruning_ratio
参数max_pruning_ratio是对所有层所应用最大剪枝比例的阈值。当剪枝器试图将单个层的所有参数移除时，此功能可以防止它这样做。
iterative_steps & iterative_sparsity_scheduler
参数iterative_steps在需要分多轮次剪枝模型的情况下非常有用。默认情况下，剪枝器会逐步增加模型的稀疏性，直到达到预期的pruning_ratio。要实现“剪枝-微调”的循环过程，可以按照以下步骤进行：
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18()
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.MagnitudeImportance(p=2)
iterative_steps = 5 # progressive pruning
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
iterative_steps=iterative_steps,
pruning_ratio=0.5, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
)
# prune the model, iteratively if necessary.
base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
for i in range(iterative_steps):
pruner.step()
macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
print("Round %d/%d, Params: %.2f M" % (i+1, iterative_steps, nparams/1e6))
# finetune your model here
# finetune(model)
# ...
print(model)
Round 1/5, Params: 9.44 M
Round 2/5, Params: 7.45 M
Round 3/5, Params: 5.71 M
Round 4/5, Params: 4.20 M
Round 5/5, Params: 2.93 M
ResNet(
(conv1): Conv2d(3, 32, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer2): Sequential(
(0): BasicBlock(
(conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(32, 64, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer3): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer4): Sequential(
(0): BasicBlock(
(conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=256, out_features=500, bias=True)
)
ignored_layers
参数ignore_layers允许您提供一个由nn.Module对象组成的列表，这些对象将在剪枝过程中被排除。需要特别注意的是，ignore_layers会冻结这些层的输出通道，输入通达则不受影响。
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18()
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.MagnitudeImportance(p=2)
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
pruning_ratio=0.5, # remove 50% channels
ignored_layers=[model.conv1, model.fc] # ignore the first & last layers
)
pruner.step()
print(model)
ResNet(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer2): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(64, 64, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer3): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer4): Sequential(
(0): BasicBlock(
(conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=256, out_features=1000, bias=True)
)
round_to
您可以将通道数量四舍五入为round_to的倍数
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18()
example_inputs = torch.randn(1, 3, 224, 224)
imp = tp.importance.MagnitudeImportance(p=2)
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs,
importance=imp,
pruning_ratio=0.3, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
round_to=10 # round to 10x. Note: 10x is not a good practice.
)
pruner.step()
print(model)
ResNet(
(conv1): Conv2d(3, 40, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer2): Sequential(
(0): BasicBlock(
(conv1): Conv2d(40, 80, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(80, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(40, 80, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(80, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(80, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer3): Sequential(
(0): BasicBlock(
(conv1): Conv2d(80, 170, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(170, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(170, 170, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(170, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(80, 170, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(170, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(170, 170, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(170, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(170, 170, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(170, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(layer4): Sequential(
(0): BasicBlock(
(conv1): Conv2d(170, 350, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(350, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(350, 350, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(350, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(downsample): Sequential(
(0): Conv2d(170, 350, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): BatchNorm2d(350, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): BasicBlock(
(conv1): Conv2d(350, 350, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(350, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(350, 350, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(350, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=350, out_features=700, bias=True)
)
Advanced Arguments
channel_groups
某些层（如nn.GroupNorm和nn.Conv2d包含group参数，这会在层内引入额外的依赖性。因此，在剪枝后需要确保所有组的大小一致。为了解决这一需求，引入了参数channel_groups，用于手动对这些通道进行分组。
pruner = tp.pruner.MagnitudePruner(
model,
example_inputs=example_inputs,
importance=importance,
iterative_steps=1,
pruning_ratio=0.5,
channel_groups = {model.group_conv1: 8} # For Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), groups=8)
)
customized_pruners
customized_pruners功能允许用户为nn.Module自定义新的剪枝器。在官方库中可以参考具体示例，该示例展示了如何为特定层定义并实现了自定义的剪枝方法。
unwrapped_parameters
在PyTorch中，某些nn.Parameter可能未封装在标准的nn.Module中，例如Transformers中cls_token和ConvNext的layer_scale。在这些情况下，可以使用unwrapped_parameters指定这些参数的剪枝维度。
root_module_types
root_module_types参数用于指定分组的“根”或者首个指定的剪枝。在许多场景中，重点在于剪枝线性层和卷积层。通过该参数可以实现针对特定层的剪枝，例如root_module_types=[nn.Conv2D, nn.Linear]。
forward_fn
此函数将在DepGraph的追踪过程中被调用。该功能适用于剪枝前需要预处理的场景。
output_transform
此函数的调用方式与forward_fn类似，主要用于对输出进行变换处理。
Frequently Asked Questions
Q: 加载剪枝后的模型失败
A: 使用 Torch-Pruning 时，模型架构会被修改，这使得原始 .py 文件中定义的结构无法兼容。因此，需要使用 torch.save(model, PATH) 保存整个模型对象，而不是仅使用 torch.save(model.state_dict(), PATH) 保存模型状态字典。这样可以确保通过 model = torch.load(PATH) 成功重新加载模型。
Q: 剪枝后 .pth 文件的大小变大
A: 剪枝后保存的 .pth 文件包含整个模型对象，因此会存储一些额外的信息。不过，实际模型的大小已经减小。要验证这一点，可以将模型加载到内存中检查，或者将剪枝后的模型导出为 ONNX 格式以确认其大小变化。
Q: 剪枝后调用 optimizer.step 报错
A: 剪枝后，需要重新创建优化器。因为剪枝后的模型参数已经被新的参数替代，如果不重新创建优化器，可能会导致优化器仍然绑定到剪枝前的旧参数。此外，注意 Torch-Pruning 无法剪枝优化器中定义的动量值（momentum）。
Q: 剪枝时出现 KeyError 错误
A: Torch-Pruning 依赖 AutoGrad 追踪计算图。因此，请确保所有层都在 forward 过程中被执行，并确保所有参数的属性 requires_grad=True。
Q: 剪枝后模型未发生变化
A: 请确认模型没有在 torch.no_grad() 的上下文中执行，同时确保所有参数的 requires_grad=True，以便网络能够正确地被追踪和剪枝。
