模型剪枝代码

一 模型剪枝的目的

1. 提高模型推理效率
2. 减少计算资源消耗

模型剪枝是指减少神经网络中的冗余参数，可以在单层、多层或整个模型中进行。分为结构化剪枝和非机构化剪枝。

二 非结构化剪枝

指在模型中直接剪掉一些不重要的参数，从而减少模型的参数量。
这种剪枝方式不考虑模型的结构，只关注参数的重要性。

1. 确定剪枝比例：首先，需要确定要剪枝的比例，即要去除的连接的比例。这可以根据模型的需求和性能要求来确定。
2. 评估连接的重要性：对于每个连接，需要评估其对模型性能的贡献。一种常见的方法是使用连接的权重作为重要性的度量。较小的权重通常表示连接不重要。
3. 剪枝连接：根据连接的重要性，选择要剪枝的连接。可以根据预先设定的阈值来确定哪些连接应该被剪枝。
4. 重新训练模型：剪枝后，需要对模型进行重新训练，以恢复模型的性能。在重新训练过程中，可以使用稀疏性正则化等技术来帮助模型恢复性能。

举个例子：

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune


# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 定义剪枝函数
def prune_model(model):
    parameters_to_prune = (
        (model.conv1, 'weight'),
        (model.conv2, 'weight'),
        (model.fc, 'weight')
    )
    # 选择了L1非结构化剪枝方法，并将剪枝比例设置为50%
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=0.5
    )


# 统计剪枝后置零的参数
def display_pruned_zeros(model):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
            mask = getattr(module, 'weight_mask', None)
            if mask is not None:
                pruned_zeros = torch.sum(mask == 0).item()
                print(f"Layer {name}: {pruned_zeros} pruned zeros")



if __name__ == '__main__':
    model = MyModel()
    print(model)
    prune_model(model)
    display_pruned_zeros(model)
    
	# 输出结果
    """
	MyModel(
	  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
	  (relu): ReLU()
	  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
	  (fc): Linear(in_features=64, out_features=10, bias=True)
	)
	Layer conv1: 200 pruned zeros
	Layer conv2: 19299 pruned zeros
	Layer fc: 117 pruned zeros
	"""

	#
	# 重新训练模型
	# ...
	
	# 使用剪枝后的模型进行推理
	# ...

三 结构化剪枝

明天了再写