3. 非结构化剪枝：权重剪枝的原理、阈值设定方法、PyTorch实现权重剪枝

好，咱们今天聊点硬核的——非结构化剪枝。说白了，就是直接把模型里那些「不重要的」权重给干掉。我刚开始接触这个方向时，总觉得这操作有点暴力，把权重置零，网络还能好好干活吗？后来发现，不仅能干，还能干得漂亮。

3.1 权重剪枝的原理：为什么能剪？

深度学习模型，尤其是大模型，其实挺「臃肿」的。你想想看，几百万甚至上亿的参数，真正起决定性作用的，可能连一半都不到。剩下的那些，要么数值极小，要么对最终输出影响微乎其微。

权重剪枝的核心思想就一句话：把绝对值小于某个阈值的权重直接干掉（置为0）。为什么是绝对值？因为权重的正负只代表对激活的抑制或兴奋，大小才代表影响力。

核心原理： 神经网络中存在大量冗余参数。通过移除这些冗余，可以在几乎不损失精度的情况下，大幅减少模型的计算量和存储量。

我在项目中遇到过这样一个情况：一个用于图像分类的ResNet-50，剪掉60%的权重后，准确率只掉了0.3%。但模型大小直接从90MB降到了36MB。嗯，这买卖划算。

3.2 阈值设定方法：剪多少才合适？

阈值怎么定？这是个好问题。定高了，模型精度崩了；定低了，剪了个寂寞。我个人习惯用以下几种方法：

3.2.1 全局阈值法

最简单粗暴。对整个模型的所有权重，设定一个统一的阈值。比如，绝对值小于0.01的权重全部干掉。

  • 优点： 实现简单，一视同仁。
  • 缺点： 不同层的权重分布差异很大。浅层权重通常较大，深层权重较小。一刀切容易把深层剪废了。

3.2.2 按比例剪枝（百分比阈值）

这是我最常用的方法。不设绝对值阈值，而是设定一个剪枝比例。比如剪掉50%的权重，那就把所有权重按绝对值排序，砍掉后50%的。

  • 优点： 每层都能保证剪掉相同比例，分布更均匀。
  • 缺点： 有些层可能本身就很重要，剪太多反而不好。

3.2.3 按层自适应阈值

更精细的做法。对每一层单独计算阈值。比如，以该层权重标准差的倍数作为阈值。

我的经验： 如果你刚开始做剪枝，建议从按比例剪枝入手。设置一个30%-50%的剪枝率，先跑一遍看看效果。我曾经一上来就剪了80%，结果模型直接崩了，损失函数飞到了天上。后来老老实实从30%开始，逐步往上加。

3.3 PyTorch实现权重剪枝：手把手带你写代码

好了，理论说完了，咱们直接上代码。PyTorch从1.4版本开始，内置了剪枝工具，在 torch.nn.utils.prune 模块里。我个人觉得这个API设计得挺优雅的。

3.3.1 最简单的全局剪枝示例

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

# 定义一个简单的网络
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleNet()

# 对 fc1 层进行 L1 非结构化剪枝，剪掉 40% 的权重
prune.l1_unstructured(model.fc1, name='weight', amount=0.4)

# 查看剪枝后的权重
print(model.fc1.weight)
# 你会发现很多权重变成了 0.0

看到了吗？就这么几行代码，40%的权重就被干掉了。但这里有个坑——剪枝后的权重变成了一个 masked tensor。什么意思？就是原来的 weight 属性被替换成了 weight_orig（原始权重）和 weight_mask（掩码）。

我曾经踩过的坑： 剪枝后直接保存模型，下次加载时发现权重还是原来的样子。为什么？因为 prune 默认是永久性移除掩码的，但如果你只保存了 state_dict，掩码信息可能丢失。正确的做法是调用 prune.remove(model.fc1, 'weight') 来固化剪枝结果。

3.3.2 全局剪枝：给整个模型动刀

有时候你想对整个模型的所有层统一剪枝比例，而不是每层单独剪。这时候用全局剪枝：

# 收集所有要剪枝的参数
parameters_to_prune = [
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'),
]

# 全局剪枝，剪掉整体 50% 的权重
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.5,
)

全局剪枝的好处是，它会自动把「不重要的层」多剪一些，「重要的层」少剪一些。说白了，就是让模型自己决定哪些层可以多牺牲一点。

3.3.3 自定义剪枝阈值

如果你不想用内置的L1剪枝，想自己定义阈值，可以这样：

class CustomThresholdPruning(prune.BasePruningMethod):
    """自定义阈值剪枝：绝对值小于 threshold 的权重全部剪掉"""
    PRUNING_TYPE = 'unstructured'
    
    def __init__(self, threshold):
        self.threshold = threshold
    
    def compute_mask(self, t, default_mask):
        # 生成掩码：绝对值大于阈值的保留，小于的干掉
        mask = torch.abs(t) > self.threshold
        return mask

# 使用自定义剪枝
CustomThresholdPruning.apply(model.fc1, 'weight', threshold=0.01)

嗯，这里要注意：compute_mask 返回的是布尔张量，True 表示保留，False 表示剪掉。这个逻辑别搞反了，我曾经就犯过这个错，结果把重要的权重全剪了，模型直接废了。

3.4 剪枝后的模型微调：别忘了一步

剪完枝，模型精度肯定会掉一点。这时候需要做一件事——微调（Fine-tuning）。用原来的训练数据，再训几个epoch，让模型适应新的稀疏结构。

我的习惯： 剪枝后先用较小的学习率（比如原学习率的1/10）微调5-10个epoch。如果精度恢复得不错，再逐渐加大学习率。千万别一上来就用大学习率，稀疏后的网络很脆弱，容易训崩。

3.5 总结一下

权重剪枝说白了就是「去芜存菁」。把那些不重要的权重干掉，模型变小了，推理变快了，精度还几乎不掉。你想想看，这多划算。

最后给新手一个建议：别追求极致的剪枝率。我见过有人为了炫技，把模型剪到只剩10%的参数，精度掉了5个点，然后花大量时间做知识蒸馏去恢复精度。说实话，性价比不高。工业界更看重的是「在精度损失可接受范围内，最大化压缩率」。一般剪掉40%-60%是比较稳妥的选择。

好了，这一章就到这里。下一章咱们聊聊结构化剪枝，那玩意儿更狠，直接干掉整个卷积核或者通道。到时候见。

← 上一章 📖 返回目录 下一章 →