模型轻量化基础：剪枝、蒸馏与量化

好，咱们进入第三章。模型轻量化，说白了就是给模型「减肥」。你想想看，BEV模型动辄几十上百兆，想在嵌入式芯片上跑起来，不瘦身怎么行？

我个人习惯把轻量化技术分成三大类：剪枝、蒸馏、量化。这三板斧用好了，模型体积能缩到原来的十分之一，速度翻几倍，精度损失却可以控制在1%以内。今天我就把这三种方法掰开揉碎了讲。

3.1 模型剪枝：结构化 vs 非结构化

剪枝，就是砍掉模型中不重要的连接或通道。我刚开始做剪枝时，总觉得这玩意儿玄乎——砍掉参数，精度不掉？后来发现，神经网络里确实有很多冗余参数，砍掉它们影响不大。

3.1.1 非结构化剪枝

非结构化剪枝，是对单个权重下手。设定一个阈值，低于阈值的权重直接置零。

# 非结构化剪枝示例
import torch
import torch.nn.utils.prune as prune

model = YourBEVModel()
# 对卷积层进行L1范数剪枝，保留30%的权重
prune.l1_unstructured(model.conv1, name='weight', amount=0.7)
# 剪枝后权重变成稀疏矩阵
print(model.conv1.weight_mask)  # 0/1掩码

我在项目中遇到过这种情况：用非结构化剪枝把模型压缩了5倍，结果在ARM CPU上跑，速度反而慢了。为什么？因为稀疏矩阵的存储和计算都不规则，CPU根本加速不了。

3.1.2 结构化剪枝

结构化剪枝就实在多了。它直接砍掉整个通道或整个卷积核。这样剪完，模型结构还是规整的，硬件友好。

# 结构化剪枝：按通道重要性排序
def channel_prune(model, prune_ratio=0.3):
    # 计算每个通道的L2范数作为重要性指标
    importance = []
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() == 4:
            # 对每个输出通道计算L2范数
            l2_norm = torch.norm(param.view(param.size(0), -1), dim=1)
            importance.append(l2_norm)
    
    # 合并所有层的重要性，设定全局阈值
    all_importance = torch.cat(importance)
    threshold = torch.quantile(all_importance, prune_ratio)
    
    # 砍掉低于阈值的通道
    # ... 实际剪枝操作需要重写网络结构

3.2 知识蒸馏：大模型教小模型

知识蒸馏，说白了就是「师夷长技以制夷」。用一个大的教师模型，去教一个小学生模型。教师模型的知识，通过软标签（soft label）传递给学生。

我记得第一次做蒸馏时，觉得这方法太巧妙了。教师模型不仅告诉学生「正确答案是什么」，还告诉它「哪些错误答案比较接近」。这种信息量，比单纯的硬标签丰富得多。

3.2.1 蒸馏的损失函数

def distillation_loss(student_logits, teacher_logits, labels, T=4, alpha=0.7):
    """
    T: 温度参数，控制软标签的平滑程度
    alpha: 软标签损失的权重
    """
    # 软标签损失：KL散度
    soft_loss = nn.KLDivLoss()(
        F.log_softmax(student_logits / T, dim=1),
        F.softmax(teacher_logits / T, dim=1)
    ) * (T * T)
    
    # 硬标签损失：交叉熵
    hard_loss = nn.CrossEntropyLoss()(student_logits, labels)
    
    # 总损失
    return alpha * soft_loss + (1 - alpha) * hard_loss

3.2.2 BEV模型蒸馏的特殊性

BEV模型的蒸馏，和图像分类不太一样。BEV输出的是特征图或检测结果，不是分类概率。所以我们需要用特征蒸馏或检测蒸馏。

# BEV特征蒸馏示例
def bev_feature_distill(student_feat, teacher_feat):
    """
    student_feat: 学生模型输出的BEV特征图 [B, C, H, W]
    teacher_feat: 教师模型输出的BEV特征图 [B, C', H, W]
    """
    # 如果通道数不一致，加个1x1卷积对齐
    if student_feat.size(1) != teacher_feat.size(1):
        student_feat = nn.Conv2d(student_feat.size(1), teacher_feat.size(1), 1)(student_feat)
    
    # 计算MSE损失
    loss = F.mse_loss(student_feat, teacher_feat.detach())
    return loss

3.3 量化：从FP32到INT8

量化，就是把模型从32位浮点数，变成8位整数。模型体积直接缩到四分之一，推理速度翻倍。嵌入式平台上的标配操作。

3.3.1 训练后量化（PTQ）

PTQ最简单，模型训练完直接量化。用一小批校准数据，统计每层激活值的范围，然后做量化。

# PyTorch PTQ示例
import torch.quantization as quant

model = YourBEVModel()
model.eval()

# 配置量化后端
model.qconfig = quant.get_default_qconfig('fbgemm')
# 准备量化
model_prepared = quant.prepare(model)
# 用校准数据跑一遍，统计激活值范围
with torch.no_grad():
    for data in calib_loader:
        model_prepared(data)
# 转换为量化模型
model_quantized = quant.convert(model_prepared)

3.3.2 量化感知训练（QAT）

QAT在训练过程中模拟量化效果，让模型适应量化带来的精度损失。效果比PTQ好，但需要重新训练。

# QAT训练流程
model = YourBEVModel()
model.train()

# 配置QAT
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
model_prepared = quant.prepare_qat(model)

# 正常训练，但前向传播时模拟量化
for epoch in range(num_epochs):
    for data, target in train_loader:
        output = model_prepared(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 训练完成后转换为量化模型
model_quantized = quant.convert(model_prepared.eval())

方法	精度损失	需要训练	适用场景
PTQ	1%-3%	否	快速部署、资源有限
QAT	0.5%-1%	是	精度要求高、有时间训练

3.4 三种方法的组合策略

在实际项目中，这三种方法不是互斥的。我常用的组合是：先剪枝，再蒸馏，最后量化。

1. 剪枝：砍掉30%-50%的通道，模型体积减半
2. 蒸馏：用原始大模型做教师，恢复剪枝后的精度损失
3. 量化：INT8量化，体积再缩四分之一，速度翻倍

这样一套组合拳下来，模型体积能压缩到原来的十分之一，精度损失控制在1%以内。我在一个BEV感知项目上就是这么干的，模型从200MB压缩到18MB，在Jetson Orin上推理速度从30ms降到8ms。