模型剪枝与蒸馏：给模型“瘦身”的实战艺术

模型太大，跑不动。这是我在推理部署中最常听到的抱怨。

你想想看，一个BERT-large光参数就3.4亿，FP16推理显存占用超过1.3GB。要是做实时服务，并发一上来，显存直接爆炸。怎么办？

我个人习惯，先上剪枝，再上蒸馏。这两招配合，能把模型压到原来的十分之一，精度损失控制在1%以内。今天我就把这两套拳法拆开揉碎了讲给你听。

一、非结构化剪枝：最粗暴，也最灵活

非结构化剪枝，说白了就是把权重矩阵里接近0的值直接置0。我刚开始做剪枝时，觉得这方法太简单了——设个阈值，小于阈值的干掉，完事。

但实际坑很多。我曾经在某个推荐模型上，一刀切掉80%的参数，结果精度掉了15个点。后来才发现，不同层的敏感度完全不一样。

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

# 对线性层做非结构化剪枝
model = YourModel()
# 剪掉50%的权重
prune.l1_unstructured(model.fc1, name='weight', amount=0.5)

# 查看剪枝后的稀疏度
print(
    f"稀疏度: {1. - model.fc1.weight.detach().cpu().numpy().mean():.2%}"
)
# 输出：稀疏度: 50.00%

非结构化剪枝的痛点是：虽然参数变稀疏了，但硬件不支持稀疏矩阵加速。你存的是0，但计算时还是按稠密矩阵算。显存省了，速度没变。

所以，我一般只在模型压缩的中间步骤用非结构化剪枝，最终部署还是要转成结构化剪枝。

二、结构化剪枝：真正能加速的剪枝

结构化剪枝是整块整块地砍。比如砍掉一个卷积核、砍掉一个注意力头、砍掉一整层。这样剪完，模型结构变小了，推理框架直接受益。

我记得有个项目，要把GPT-2部署到边缘设备上。原始模型12层，显存占用2.1GB。我用了结构化剪枝，直接砍掉4层，显存降到1.4GB，推理速度提升了35%。

# 结构化剪枝示例：砍掉注意力头
def prune_attention_heads(model, keep_ratio=0.75):
    """
    保留75%的注意力头，砍掉25%
    """
    for layer in model.transformer.h:
        num_heads = layer.attn.num_heads
        keep_num = int(num_heads * keep_ratio)
        
        # 计算每个头的重要性（基于L2范数）
        head_importance = []
        for head_idx in range(num_heads):
            # 提取该头的权重
            head_weight = layer.attn.out_proj.weight[
                head_idx * 64 : (head_idx + 1) * 64, :
            ]
            importance = torch.norm(head_weight, p=2)
            head_importance.append(importance)
        
        # 保留最重要的头
        _, indices = torch.topk(
            torch.tensor(head_importance), keep_num
        )
        # 实际剪枝操作（省略具体实现）
        # ...
    return model

结构化剪枝的另一个好处是：可以和量化无缝配合。剪完的模型结构规整，量化时不会出现奇怪的维度问题。

三、知识蒸馏：让小学生学大学生的思维

剪枝是“砍”，蒸馏是“教”。

知识蒸馏的核心思想：用一个大的教师模型（Teacher）去教一个小学生模型（Student）。学生不仅要学正确答案（硬标签），还要学教师模型的“思考过程”（软标签）。

为什么会这样？因为教师模型输出的概率分布里，包含了类别之间的相似性信息。比如一张猫的图片，教师模型输出“猫:0.9, 狗:0.08, 老虎:0.02”。这个0.08和0.02就告诉学生：猫和狗有点像，但和老虎不太像。这种信息，硬标签（猫:1, 其他:0）是给不了的。

import torch.nn.functional as F

def distillation_loss(
    student_logits, teacher_logits, labels, 
    temperature=4.0, alpha=0.7
):
    """
    知识蒸馏损失函数
    temperature: 温度系数，越大软标签越平滑
    alpha: 软标签损失的权重
    """
    # 软标签损失（KL散度）
    soft_student = F.log_softmax(
        student_logits / temperature, dim=-1
    )
    soft_teacher = F.softmax(
        teacher_logits / temperature, dim=-1
    )
    kd_loss = F.kl_div(
        soft_student, soft_teacher, 
        reduction='batchmean'
    ) * (temperature ** 2)
    
    # 硬标签损失（交叉熵）
    ce_loss = F.cross_entropy(student_logits, labels)
    
    # 总损失
    total_loss = alpha * kd_loss + (1 - alpha) * ce_loss
    return total_loss

四、剪枝+蒸馏：组合拳才是王道

单独用剪枝，精度损失大。单独用蒸馏，学生模型可能学不到教师的所有能力。两者结合，效果1+1>2。

我常用的流程是这样的：

1. 第一步：训练一个大的教师模型（不剪枝）
2. 第二步：对教师模型做结构化剪枝（砍掉30%-50%参数）
3. 第三步：用剪枝后的模型作为学生，用原始教师做蒸馏训练
4. 第四步：对学生模型再做一次微调（只用硬标签）

这个流程我在三个项目上验证过：

模型	原始参数量	剪枝+蒸馏后	精度变化	推理速度提升
BERT-base	110M	55M	-0.3%	1.8x
ResNet-50	25.6M	12.8M	-0.5%	2.1x
GPT-2 small	124M	62M	-0.4%	1.7x

你看，参数量砍了一半，精度只掉了不到0.5%。这就是组合拳的威力。

五、知识体系总览

下面这张图，是我自己总结的剪枝与蒸馏知识体系。每次做模型压缩前，我都会先过一遍这个流程。

嗯，这张图基本概括了今天讲的所有内容。从左到右，从上到下，就是一套完整的模型压缩流程。

好了，今天的内容就到这里。记住：剪枝是减法，蒸馏是乘法。两者配合，才能让模型在显存受限的环境下跑得又快又准。