模型轻量化技术：让AI模型在嵌入式设备上跑起来

做嵌入式AI部署，最头疼的问题是什么？

模型太大，跑不动。算力不够，内存吃紧。我刚开始接触这个领域时，总觉得把训练好的模型直接丢到芯片上就能用。结果呢？一个几十MB的模型，在Cortex-M上根本跑不起来，连加载都加载不完。

所以，模型轻量化不是「锦上添花」，而是「生死攸关」。今天我就把这几年积累的轻量化经验，掰开揉碎了讲给你听。

模型剪枝：把「赘肉」剪掉

模型剪枝，说白了就是去掉那些不重要的参数。神经网络里有很多权重其实接近于零，留着它们只会浪费计算资源。

非结构化剪枝

这种方式比较「暴力」——直接把小权重置为零。我习惯用阈值法：设定一个阈值，低于这个值的权重全部干掉。

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

# 非结构化剪枝示例
model = YourModel()
# 对卷积层进行剪枝，保留20%的权重
prune.l1_unstructured(model.conv1, name='weight', amount=0.8)
# 剪枝后需要重新训练恢复精度

结构化剪枝

这个就聪明多了——直接剪掉整个通道或整个卷积核。硬件友好，加速效果明显。

# 结构化剪枝：按通道剪枝
# 计算每个通道的L1范数，保留重要的通道
def channel_prune(model, prune_ratio=0.3):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            # 计算每个卷积核的L1范数
            l1_norm = torch.norm(module.weight.data, p=1, dim=(1,2,3))
            # 找到要剪掉的通道索引
            num_prune = int(module.out_channels * prune_ratio)
            _, indices = torch.topk(l1_norm, num_prune, largest=False)
            # 这里需要实际修改网络结构
            # 具体实现略复杂，建议用torch-pruning库

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

量化，就是把FP32的浮点数变成INT8甚至INT4。你想想看，一个参数从32位变成8位，模型直接缩小4倍。这买卖划算吧？

INT8量化

这是目前最成熟、应用最广的方案。TensorRT、ONNX Runtime、TFLite都支持得很好。

# PyTorch量化示例
import torch.quantization as quant

model = YourModel()
model.eval()

# 配置量化
model.qconfig = quant.get_default_qconfig('fbgemm')
model = quant.prepare(model, inplace=True)

# 校准：用少量数据跑一遍，收集激活值的范围
calibrate(model, calib_loader)

# 转换量化模型
model = quant.convert(model, inplace=True)

量化类型	模型大小缩减	精度损失	硬件要求
FP32	1x	基准	通用
INT8	4x	<1%	需要INT8加速单元
INT4	8x	1-3%	部分芯片支持

INT4量化

这个就比较激进了。我曾在某个项目中尝试INT4量化，模型小了8倍，但精度掉了3%。后来通过量化感知训练（QAT）才把精度拉回来。

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

知识蒸馏的思路很有意思——用一个大的教师模型去训练一个小学生模型。教师模型知道的知识，通过软标签（soft label）传递给学生。

# 知识蒸馏损失函数
def distillation_loss(student_logits, teacher_logits, labels, T=4, alpha=0.7):
    # 软标签损失：让学生模仿教师的输出分布
    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

温度T是个超参数。T越大，软标签越平滑，学生学到的东西越「泛化」。我一般设T=4，alpha=0.7，效果比较稳。

轻量化网络设计：从源头解决问题

与其在训练好的模型上做剪枝量化，不如一开始就设计轻量化的网络结构。MobileNet、ShuffleNet、EfficientNet-Lite，这三兄弟是嵌入式AI的「三驾马车」。

MobileNetV3

深度可分离卷积 + 倒残差结构 + SE注意力。我特别喜欢它的h-swish激活函数，计算量小，效果还不赖。

# MobileNetV3的核心模块
class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super().__init__()
        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = stride == 1 and inp == oup
        
        layers = []
        if expand_ratio != 1:
            # 1x1升维
            layers.append(nn.Conv2d(inp, hidden_dim, 1, bias=False))
            layers.append(nn.BatchNorm2d(hidden_dim))
            layers.append(nn.Hardswish())
        # 3x3深度可分离卷积
        layers.extend([
            nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, 
                     groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.Hardswish(),
            # 1x1降维
            nn.Conv2d(hidden_dim, oup, 1, bias=False),
            nn.BatchNorm2d(oup),
        ])
        self.conv = nn.Sequential(*layers)

ShuffleNetV2

这个网络的设计原则很实在——考虑实际硬件特性。通道混洗（channel shuffle）操作让信息在不同组之间流动，计算量小，精度高。

EfficientNet-Lite

EfficientNet的嵌入式版本。去掉了SE模块（嵌入式上跑不动），改用ReLU6激活函数。我实测过，在树莓派上比MobileNetV3快15%左右。

组合拳：多种技术一起上

实际项目中，我很少只用一种技术。通常是先设计轻量化网络，然后做结构化剪枝，再INT8量化，最后用知识蒸馏微调。这一套组合拳打下来，模型能缩小10-20倍，精度损失控制在1%以内。

举个例子，我去年做的一个视觉检测项目：

1. 基础模型：MobileNetV3-Large，4.2MB
2. 结构化剪枝30%：降到2.9MB，精度掉0.3%
3. INT8量化：降到0.73MB，精度再掉0.5%
4. 知识蒸馏微调：精度恢复0.4%

最终模型0.73MB，精度只比原始模型低0.4%。在STM32H7上跑到了30fps，完美满足需求。