今天开始学剪枝

剪枝技术：提升模型效率的策略

最新推荐文章于 2025-12-03 22:01:18 发布

原创最新推荐文章于 2025-12-03 22:01:18 发布 · 969 阅读

24 ·

CC 4.0 BY-SA版权

文章标签：

#剪枝

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import prune

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleNet()

# 随机剪枝
prune.random_unstructured(model.fc1, name='weight', amount=0.8)

# 检查剪枝状态
print(prune.is_pruned(model.fc1))

# 获取剪枝掩码
mask = model.fc1.weight_mask
print(mask)

# 移除剪枝
prune.remove(model.fc1, name='weight')

一个简单的剪枝示例，这里剪掉了80%是为了演示效果，

实际操作中可不能这么狠心，尤其对于这种随机剪枝

输出结果：

True
tensor([[1., 0., 1.,  ..., 0., 0., 0.],
        [0., 0., 1.,  ..., 0., 0., 1.],
        [1., 0., 0.,  ..., 0., 1., 0.],
        ...,
        [0., 0., 0.,  ..., 1., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 1.],
        [0., 0., 0.,  ..., 0., 1., 0.]])

Linear(in_features=784, out_features=256, bias=True)

再来看看几种常见的剪枝

torch.nn.utils.prune 提供了多种剪枝方法，常用的有以下几种：

1. L1 剪枝（L1 Unstructured Pruning）

描述：基于权重的 L1 范数剪枝，移除权重最小的部分。

用法：

python

复制

prune.l1_unstructured(module, name='weight', amount=0.5)

2. 随机剪枝（Random Unstructured Pruning）

描述：随机移除权重。

用法：

python

复制

prune.random_unstructured(module, name='weight', amount=0.5)

3. L2 结构化剪枝（L2 Structured Pruning）

描述：基于权重的 L2 范数剪枝，移除整个神经元或通道。

用法：

python

复制

prune.ln_structured(module, name='weight', amount=0.5, n=2, dim=0)

4. 全局剪枝（Global Pruning）

描述：跨多个层进行全局剪枝。

用法：

python

复制

parameters_to_prune = ((model.fc1, 'weight'), (model.fc2, 'weight'))
prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.5)

5. 自定义剪枝（Custom Pruning）

描述：自定义剪枝方法。

用法：

python

复制

def custom_pruning_mask(weight, amount):
    mask = torch.ones_like(weight)
    # 自定义剪枝逻辑
    return mask

prune.custom_from_mask(module, name='weight', mask=custom_pruning_mask(module.weight, amount=0.5))

6. 移除剪枝（Remove Pruning）

描述：移除剪枝，恢复原始权重。
用法：

python

复制
```
prune.remove(module, name='weight')
```

7. 检查剪枝状态（Check Pruning Status）

描述：检查模块是否被剪枝。
用法：

python

复制
```
print(prune.is_pruned(module))
```

8. 获取剪枝掩码（Get Pruning Mask）

描述：获取剪枝掩码。

用法：

python

复制

mask = prune.get_mask(module, name='weight')

示例代码

python

复制

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import prune

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleNet()

# L1 剪枝
prune.l1_unstructured(model.fc1, name='weight', amount=0.5)

# 检查剪枝状态
print(prune.is_pruned(model.fc1))

# 获取剪枝掩码
mask = prune.get_mask(model.fc1, name='weight')

# 移除剪枝
prune.remove(model.fc1, name='weight')