3. 模型轻量化技术（上）：模型剪枝原理与实现

模型剪枝，说白了就是给神经网络“瘦身”。

你想想看，一个训练好的模型，里面很多参数其实都是冗余的。就像一棵树，有些枝叶不结果，剪掉反而能让主干长得更好。我刚开始做端侧部署时，第一个遇到的坎就是模型太大，MCU根本塞不下。后来发现，剪枝是最直接、最有效的轻量化手段之一。

3.1 剪枝的核心思想

剪枝的本质，是移除模型中不重要的权重或神经元。判断“重要性”的标准有很多，最常见的是看权重绝对值的大小。绝对值小的权重，对最终输出的贡献就小，剪掉它们对精度影响也最小。

我在项目中遇到过这样一个情况：一个3MB的模型，剪掉60%的参数后，精度只掉了0.5%。但模型大小直接缩到了1.2MB，推理速度也快了近一倍。嗯，这种性价比，谁用谁知道。

3.2 结构化剪枝 vs 非结构化剪枝

这是剪枝领域最基础、也最重要的分类。我建议你先把这两个概念吃透，后面用工具时才不会懵。

对比维度	结构化剪枝	非结构化剪枝
剪枝粒度	以通道（Channel）、卷积核（Filter）或层（Layer）为单位	以单个权重（Weight）为单位
硬件加速	直接支持，无需特殊硬件	需要稀疏计算库或专用硬件
压缩效果	模型体积减小，推理速度明显提升	模型体积减小，但推理速度可能不变
精度损失	相对较大，需要仔细调参	相对较小，保留更多细节
部署难度	低，直接导出即可	高，需要稀疏矩阵支持

举个例子。结构化剪枝就像把一整块木板锯掉一截，尺寸变了，但结构完整。非结构化剪枝则像在木板上打孔，孔多了木板变轻，但形状没变。

我个人习惯在MCU部署时优先用结构化剪枝。为什么？因为Cortex-M系列内核没有稀疏矩阵加速指令，非结构化剪枝带来的稀疏权重，反而会让内存访问变得碎片化，速度不升反降。

3.3 剪枝的完整流程

一个标准的剪枝流程，通常包含以下几步。我把它画成了流程图，方便你理解。

3.4 基于TensorFlow的剪枝工具

TensorFlow提供了tensorflow-model-optimization库，里面集成了剪枝API。我个人觉得它的设计很清晰，适合快速上手。

下面是一个结构化剪枝的示例。注意，这里用的是多项式衰减调度策略，剪枝率从0逐步增加到目标值。

import tensorflow as tf
from tensorflow_model_optimization.sparsity import keras as sparsity

# 定义剪枝参数
pruning_params = {
    'pruning_schedule': sparsity.PolynomialDecay(
        initial_sparsity=0.0,      # 初始剪枝率
        final_sparsity=0.5,        # 目标剪枝率（剪掉50%）
        begin_step=0,
        end_step=1000,
        frequency=100
    )
}

# 对模型应用剪枝
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(32, 32, 3)),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

pruned_model = sparsity.prune_low_magnitude(model, **pruning_params)

# 编译并训练
pruned_model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 训练时需要使用回调
callbacks = [
    sparsity.UpdatePruningStep(),
    sparsity.PruningSummaries(log_dir='./logs')
]

pruned_model.fit(x_train, y_train, epochs=5, callbacks=callbacks)

# 导出剪枝后的模型
final_model = sparsity.strip_pruning(pruned_model)
final_model.save('pruned_model.h5')

3.5 基于PyTorch的剪枝工具

PyTorch的torch.nn.utils.prune模块提供了更灵活的剪枝方式。它支持全局剪枝和局部剪枝，你可以精确控制每一层的剪枝率。

下面是一个非结构化剪枝的例子。不过，我建议你在MCU部署时，把它改成结构化剪枝。

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

# 定义一个简单模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.fc = nn.Linear(32 * 30 * 30, 10)

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

model = SimpleCNN()

# 对conv1层进行非结构化剪枝（剪掉30%的权重）
prune.l1_unstructured(model.conv1, name='weight', amount=0.3)

# 对fc层进行结构化剪枝（剪掉20%的神经元）
prune.ln_structured(model.fc, name='weight', amount=0.2, n=2, dim=0)

# 查看剪枝后的权重掩码
print(model.conv1.weight_mask)

# 永久化剪枝（移除掩码，只保留剪枝后的权重）
prune.remove(model.conv1, 'weight')
prune.remove(model.fc, 'weight')

3.6 剪枝率的选择策略

剪枝率不是越高越好。我总结了一个经验表格，你可以作为参考起点。

应用场景	推荐剪枝率	预期精度损失	备注
简单分类（CIFAR-10）	50% - 70%	< 1%	模型冗余度高，可大胆剪
目标检测（Tiny YOLO）	30% - 50%	1% - 3%	检测任务对细节敏感
语义分割	20% - 40%	2% - 5%	逐像素任务，谨慎剪枝
MCU端侧部署	40% - 60%	1% - 2%	结合量化效果更佳

我个人习惯从30%的剪枝率开始试。如果精度损失在可接受范围内，再逐步增加。每次增加10%，直到精度出现明显下降，然后回退到上一个值。

3.7 剪枝后的模型验证

剪枝完成后，别急着部署。先做三件事：

1. 精度验证：在验证集上跑一遍，确认精度损失在预期范围内。
2. 模型大小检查：对比剪枝前后的模型文件大小，确认压缩效果。
3. 推理速度测试：在目标硬件上跑一次推理，确认速度提升。