模型剪枝知识点整理

模型剪枝知识点整理

剪枝是深度学习模型优化的两种常见技术，用于减少模型复杂度和提升推理速度，适用于资源受限的环境。

剪枝（Pruning）

剪枝是一种通过移除模型中不重要或冗余的参数来减少模型大小和计算量的方法。剪枝通常分为以下几种类型：

1. 权重剪枝（Weight Pruning）

权重剪枝是通过移除权重矩阵中接近零的元素来减少模型的参数数量。常见的方法有：

• 非结构化剪枝（Unstructured Pruning）：逐个移除权重矩阵中的小权重。
• 结构化剪枝（Structured Pruning）：按特定结构（如整行或整列）移除权重。

示例：

import torch

# 假设有一个全连接层
fc = torch.nn.Linear(100, 100)

# 获取权重矩阵
weights = fc.weight.data.abs()

# 设定剪枝阈值
threshold = 0.01

# 应用剪枝
mask = weights > threshold
fc.weight.data *= mask

2. 通道剪枝（Channel Pruning）

通道剪枝主要用于卷积神经网络，通过移除卷积层中不重要的通道来减少计算量。常见的方法有：

• 基于重要性评分：计算每个通道的重要性分数，移除分数较低的通道。
• 基于稀疏性：通过增加稀疏正则项，训练过程中自然使某些通道稀疏，再进行剪枝。

import torch
import torch.nn as nn

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

model = ConvNet()

# 获取卷积层的权重
weights = model.conv1.weight.data.abs()

# 计算每个通道的L1范数
channel_importance = torch.sum(weights, dim=[1, 2, 3])

# 设定剪枝阈值
threshold = torch.topk(channel_importance, k=32, largest=True).values[-1]

# 应用剪枝
mask = channel_importance > threshold
model.conv1.weight.data *= mask.view(-1, 1, 1, 1)

3. 层剪枝（Layer Pruning）

层剪枝是移除整个网络层，以减少模型的计算深度。这种方法较为激进，通常结合模型架构搜索（NAS）使用。

import torch.nn as nn

class LayerPrunedNet(nn.Module):
    def __init__(self, use_layer=True):
        super(LayerPrunedNet, self).__init__()
        self.use_layer = use_layer
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
    
    def forward(self, x):
        x = self.conv1(x)
        if self.use_layer:
            x = self.conv2(x)
        return x

# 初始化网络，选择是否使用第二层
model = LayerPrunedNet(use_layer=False)