模型轻量化技术（上）：模型剪枝、权重量化与知识蒸馏

各位同学，今天我们来聊聊嵌入式AI里最实在的话题——模型轻量化。

说实话，我刚开始做嵌入式AI那会儿，踩过不少坑。最典型的一次，我在服务器上训了一个ResNet-50，准确率92%，心里美滋滋。结果往STM32上一部署，好家伙，Flash直接爆了，RAM也不够用。那感觉，就像你做好了满汉全席，结果发现碗太小装不下。

从那以后，我学乖了。模型轻量化不是可选项，而是必选项。今天咱们就聊三个最核心的技术：模型剪枝、权重量化、知识蒸馏。这三板斧用好了，模型体积能缩到原来的十分之一，速度还能翻倍。

一、模型剪枝：给神经网络"瘦身"

模型剪枝，说白了就是去掉那些不重要的连接或神经元。你想想看，一个神经网络里，很多权重其实接近于零，它们对最终结果贡献很小。去掉它们，模型照样能工作。

我个人习惯把剪枝分成两类：非结构化和结构化。

1.1 非结构化剪枝

非结构化剪枝，就是逐个权重地剪。设定一个阈值，低于这个阈值的权重直接置零。好处是压缩率高，坏处是——嗯，剪完后的权重矩阵变得稀疏了，普通硬件跑起来反而不快。

我在项目中遇到过这种情况：用PyTorch做了非结构化剪枝，模型大小从50MB降到了15MB，心里挺高兴。结果部署到ARM Cortex-M上，推理速度反而慢了30%。为什么？因为稀疏矩阵在通用CPU上没法高效计算，你得用专门的稀疏计算库才行。

1.2 结构化剪枝

结构化剪枝，剪的是整个通道、整个卷积核或者整个层。剪完之后网络结构还是规整的，硬件跑起来很友好。

举个例子，一个卷积层有64个输出通道，结构化剪枝可能只保留32个。这样一来，计算量直接减半，而且不需要特殊硬件支持。

# PyTorch结构化剪枝示例（剪掉BN层中gamma值小的通道）
import torch
import torch.nn.utils.prune as prune

model = torchvision.models.resnet18(pretrained=True)

# 对第一个卷积层做结构化剪枝，保留50%通道
prune.ln_structured(
    model.conv1, 
    name='weight', 
    amount=0.5, 
    n=2,  # L2范数
    dim=0  # 按输出通道剪
)

# 剪枝后需要重新打包参数
prune.remove(model.conv1, 'weight')
print(f"剪枝后卷积核形状: {model.conv1.weight.shape}")
# 输出: torch.Size([32, 3, 7, 7])  — 通道数从64减到32

二、权重量化：用更少的比特表示权重

量化，说白了就是把模型里的浮点数（FP32）换成整数（INT8或INT4）。FP32一个数占4字节，INT8只占1字节，INT4更狠，半个字节。模型体积直接缩到原来的1/4甚至1/8。

但量化不是简单的四舍五入。你得考虑数值范围、精度损失、量化误差补偿等问题。

2.1 INT8量化

INT8量化是目前最成熟、应用最广的方案。它把浮点数映射到[-128, 127]的整数范围。映射方式有两种：对称量化和非对称量化。

量化类型	映射范围	适用场景	精度损失
对称量化	[-128, 127]	权重（通常对称分布）	较小
非对称量化	[0, 255]	激活值（通常非对称分布）	更小

我建议你直接用TensorRT或ONNX Runtime的量化工具，它们会自动选择最优的量化方案。自己手写量化校准，说实话，容易翻车。

# 使用ONNX Runtime做INT8量化
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

model_fp32 = 'resnet18.onnx'
model_int8 = 'resnet18_int8.onnx'

# 动态量化：只量化权重，激活值还是浮点
quantize_dynamic(
    model_fp32,
    model_int8,
    weight_type=QuantType.QInt8
)

print("INT8量化完成！模型体积缩小约4倍")

2.2 INT4量化

INT4量化更激进，精度损失也更大。我一般只在资源极度受限的场景下用，比如Flash只有256KB的MCU。

INT4量化有个坑：很多硬件不支持INT4的矩阵乘法。你得用两个INT4拼成一个INT8来算，或者用查表法。嗯，这里要注意，查表法虽然快，但会占用额外的RAM。

三、知识蒸馏：大模型教小模型

知识蒸馏的思路很有意思——用一个大的、精度高的教师模型，去教一个小的、轻量的学生模型。学生模型学到的不是硬标签（比如"猫"或"狗"），而是教师模型输出的软概率分布。

为什么软标签更好？因为软标签包含了类别之间的相似性信息。比如教师模型看到一张图片，输出"猫"的概率是0.7，"狗"是0.2，"老虎"是0.1。这个分布告诉学生：猫和狗有点像，但猫和老虎更像。这些信息硬标签里是没有的。

3.1 蒸馏的损失函数

知识蒸馏的损失函数由两部分组成：

• 硬损失：学生模型输出与真实标签的交叉熵
• 软损失：学生模型输出与教师模型输出的KL散度

软损失里有个关键参数——温度T。温度越高，软标签分布越平滑，信息越丰富。我一般把T设在3-10之间，具体看任务。

# 知识蒸馏训练代码框架
import torch.nn.functional as F

def distillation_loss(student_output, teacher_output, labels, T=4, alpha=0.7):
    """
    student_output: 学生模型输出 (logits)
    teacher_output: 教师模型输出 (logits)
    labels: 真实标签
    T: 温度参数
    alpha: 软损失权重
    """
    # 硬损失
    hard_loss = F.cross_entropy(student_output, labels)
    
    # 软损失（用温度软化）
    soft_student = F.log_softmax(student_output / T, dim=1)
    soft_teacher = F.softmax(teacher_output / T, dim=1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T ** 2)
    
    # 总损失
    total_loss = alpha * soft_loss + (1 - alpha) * hard_loss
    return total_loss

3.2 实战中的经验

我做知识蒸馏时，发现几个关键点：

• 教师模型不能太强：如果教师模型准确率99%，学生模型学起来会很吃力。我一般选比学生强10-15个点的教师。
• 温度要逐步降低：训练初期用高温度（T=10），让学生多学软信息；后期降低温度（T=2），让学生更关注硬标签。
• 学生模型结构要合理：不是越小的模型越好。我试过把ResNet-50蒸馏到MobileNetV2，效果不错；但蒸馏到只有3层卷积的小网络，效果就很差。

好了，今天的内容就到这里。模型轻量化是个实践性很强的工作，光看理论没用，得动手试。我建议你找个现成的模型，按今天讲的流程走一遍，踩踩坑，收获会更大。