我要开一个超市,超市里面卖的是“图像分类、语义分割、模型量化剪枝与压缩-遇到了第一个挑战(1)”
8. 数据准备:从标注到数据集管理
在超市项目中,我们面临第一个挑战:数据标注的规模和质量。我们需要两种标注:目标检测的边界框和语义分割的像素级标签。
8.1 标注规范与工具
- 检测标注:用
labelImg框出每个商品,类别包括“可口可乐”、“乐事薯片”等数百种。注意边界框要贴合商品边缘,避免过多背景。 - 分割标注:用
labelme绘制多边形,区分“商品”、“货架板”、“空隙”、“顾客手部”等。分割标注更耗时,所以我们只对部分关键帧进行分割标注,用于训练UNet。
数据格式:检测数据转为COCO格式(一个JSON文件包含images、annotations、categories),分割数据转为标准的分割掩码(单通道PNG,像素值对应类别ID)。
8.2 数据增强与预处理
在训练YOLO和UNet时,数据增强至关重要。我们通过源码自定义增强策略:
- YOLO的数据增强:YOLOv5源码中的
datasets.py实现了Mosaic、MixUp、随机仿射变换等。Mosaic将四张图拼成一张,丰富背景和小物体。但Mosaic产生的图像分布与真实场景有差异,我们只在训练前期使用,后期关闭。1
2
3
4
5
6# YOLOv5 datasets.py 中的 LoadImagesAndLabels 类
if self.mosaic:
img, labels = load_mosaic(self, index)
# 之后可能还有 MixUp
if self.mixup:
img, labels = mixup(img, labels, *load_mosaic(self, random_index)) - UNet的数据增强:mmsegmentation的数据流水线在配置文件中定义,例如:通过修改配置,可以轻松加入新增强,如随机亮度对比度。
1
2
3
4
5
6
7
8
9
10train_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='LoadAnnotations'),
dict(type='RandomFlip', prob=0.5),
dict(type='RandomRotate', degree=10),
dict(type='PhotoMetricDistortion'),
dict(type='Normalize', **img_norm_cfg),
dict(type='DefaultFormatBundle'),
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
8.3 长尾分布处理
超市商品销量差异大,导致数据长尾。我们采用类别重采样和损失函数加权:
- 重采样:在YOLO训练中,对每张图根据其包含的类别计算采样权重,稀有类别图片被更频繁采样。
- 损失加权:在mmsegmentation中,为
CrossEntropyLoss设置class_weight参数,根据类别频率反比计算权重。
9. YOLO检测:深入正负样本分配与损失计算
前面已经提到YOLOv5的build_targets函数,但我们需要更深入理解它的源码实现,以及如何针对超市商品特点进行优化。
9.1 build_targets源码解剖
在utils/loss.py中,build_targets的输入是当前batch的预测特征(三个尺度)和所有GT。它的核心步骤:
- 生成anchor索引:每个尺度有3个anchor,共9个。对每个GT,计算它和所有anchor的宽高比
r = max(gt_w/anchor_w, anchor_w/gt_w),若r < anchor_t(默认4),则该anchor匹配。 - 选择最近的grid:取GT中心点坐标,向下取整得到grid坐标
(gi, gj)。然后考虑左上、右上、左下、右下四个邻域(最多三个),只要偏移在0.5以内就加入。 - 生成target张量:最终生成一个
(nt, 6)的张量,每行是(图像索引, 类别, 中心x, 中心y, 宽, 高),以及对应的anchor索引、grid坐标等辅助信息。
我们遇到的问题:超市商品中有些长条形的(如牙膏),宽高比可能大于4,导致没有anchor匹配。解决方案是调大anchor_t到6,并重新聚类anchor尺寸。
9.2 损失函数:CIoU与Focal Loss
YOLOv5默认用CIoU Loss,但我们发现对于小物体(如口香糖),CIoU对小偏移不敏感。我们改为EIoU Loss,它直接惩罚中心点距离和宽高差,收敛更快。
损失函数在loss.py的ComputeLoss类中,__call__方法里分别计算三个损失:
1 | |
注意分类和置信度损失用的是BCE,但正负样本严重不平衡,所以需要设置权重。YOLOv5在代码里通过tobj中正样本为1,负样本为0,但没有显式的focal loss。我们可以修改为Focal Loss:
1 | |
这能有效抑制易分类负样本的梯度。
9.3 NMS优化
推理时NMS(非极大抑制)可能成为瓶颈。YOLOv5源码用torchvision.ops.nms,但我们可以用更快的cluster NMS或DIoU NMS。在部署时,ncnn支持内置的NMS层,我们可以将NMS封装进模型,避免后处理开销。
10. 语义分割:FCN与UNet的源码实现对比
10.1 FCN源码要点
FCN的核心是将分类网络(如VGG)的全连接层替换为卷积,并添加转置卷积上采样。在mmsegmentation中,FCN的decode_head是FCNHead,其结构:
- 一个或多个卷积层(如两个3x3卷积),将backbone输出通道降至num_classes。
- 然后一个转置卷积上采样到原图尺寸(或者用双线性插值+卷积)。
源码位置:mmseg/models/decode_heads/fcn_head.py。
关键点是上采样方式。FCN通常使用转置卷积(也称反卷积),但转置卷积容易产生棋盘效应。我们改用双线性插值+1x1卷积,效果更好,且部署时ncnn支持更友好。
10.2 UNet的跳跃连接细节
UNet的decoder每个阶段先上采样,然后与encoder对应层输出拼接,再卷积。在mmsegmentation的unet_head.py中,_forward方法:
1 | |
这里reduce_conv是1x1卷积,用于将skip连接的通道数对齐到当前decoder的输入通道数。我们需要注意,如果encoder输出的空间尺寸与上采样后不完全一致(比如因padding导致),需要裁剪或插值对齐。
10.3 损失函数组合
对于分割任务,除了交叉熵,我们还可以添加Dice Loss处理类别不平衡。mmsegmentation允许组合多个损失:
1 | |
Dice Loss的计算在mmseg/models/losses/dice_loss.py,通过计算预测概率图与GT的Dice系数实现。
11. mmsegmentation高级特性:自定义数据集与模型
11.1 自定义数据集
超市分割数据不是标准格式,我们需要实现一个自定义数据集类。mmsegmentation的数据集注册机制让我们可以轻松添加:
1 | |
然后在配置文件中使用type='SupermarketDataset'。
11.2 测试时增强(TTA)
为了提升分割精度,我们可以在推理时使用TTA:对图像进行水平翻转、多尺度缩放,然后融合结果。mmsegmentation的BaseSegmentor支持aug_test方法,我们只需在配置中启用:
1 | |
源码中,aug_test会遍历所有变换,对结果取平均。
12. mmdeploy深入:模型导出与后端适配
12.1 自定义重写规则
在导出YOLOv5时,我们遇到了Focus层不被ONNX支持的问题。mmdeploy通过重写规则解决:在mmdeploy/rewriter/yolov5.py中,定义了focus_rewriter,将Focus模块替换为Slice+Concat的组合:
1 | |
这样导出的ONNX就是标准的Slice+Concat节点。
12.2 ONNX算子兼容性检查
导出后,我们使用onnx.checker.check_model验证模型,并用onnxsim简化。但有时简化后模型可能丢失动态轴信息。我们编写脚本自动化处理:
1 | |
12.3 ncnn部署的精度对齐
转换到ncnn后,我们需要验证推理精度是否与PyTorch一致。常见问题:
- 上采样方式:ncnn的
Interp层默认使用双线性插值,但PyTorch的F.interpolate有多种模式(bilinear, nearest),需对齐。 - padding策略:ncnn的卷积默认padding为0,而PyTorch可能使用其他值,需在转换时指定。
- BN层融合:ncnn的
optimize工具会自动融合Conv+BN,但需确保融合后的参数正确。我们对比了融合前后的输出,发现有些层因epsilon差异导致微小误差,于是手动调整epsilon。
12.4 多后端支持
虽然项目主要用ncnn,但我们也尝试了TensorRT以获得更高性能。mmdeploy支持TensorRT导出,但需要处理插件。例如YOLO的SiLU激活在TensorRT中需要转换为自定义插件,或者替换为ReLU。我们最终选择替换为ReLU并微调,避免插件依赖。
13. mmrazor全面应用:剪枝、量化、蒸馏的联动
13.1 结构化剪枝的细节
我们使用mmrazor的BNScalePruner对YOLO进行通道剪枝。该剪枝器在训练时对BN层的gamma添加L1正则,然后根据gamma大小生成剪枝mask。
源码关键点:
BNScalePruner继承自BasePruner,在train_step中计算正则损失,并调用update_mask更新mask。- 剪枝后,通过
get_pruned_model得到新模型。这一步会遍历所有可剪单元(MutableChannelUnit),根据mask裁剪卷积层的输入输出通道,并重新构建网络。 - 裁剪后,权重映射由
ModelConnector完成。例如,对于卷积层,它根据输入输出mask从原权重中选取对应的输入输出通道,重新排列。
我们遇到一个坑:跳跃连接两端的通道数必须一致。如果跳跃连接的一层被剪枝,另一层必须同步剪枝。mmrazor通过ChannelUnit管理依赖,可以配置input_units和output_units来保持一致性。
13.2 量化感知训练的细节
我们使用mmrazor的QATQuantizer对UNet进行量化。量化器会在模型中插入FakeQuantize模块,这些模块模拟量化误差。
源码细节:
QATQuantizer在prepare阶段遍历模型,将指定算子(如Conv2d)替换为QATConv2d。QATConv2d内部包含一个普通的Conv2d和两个FakeQuantize(分别用于输入和输出)。FakeQuantize的前向:对输入进行量化(计算scale和zero point,进行线性量化),然后反量化回浮点,保持梯度。- 量化参数的统计:可以使用
Observer(如MinMaxObserver)在训练中收集激活的min/max,更新scale。
我们配置了量化器使用对称量化(symmetric)和per-tensor scale,以兼容ncnn的INT8推理。
13.3 知识蒸馏的集成
为了提升剪枝后YOLO的精度,我们采用大模型(YOLOv5x)蒸馏小模型(YOLOv5s剪枝版)。mmrazor的蒸馏模块ConfigurableDistiller可以指定蒸馏损失:
1 | |
这里KLDivergence是logit蒸馏,FeatureLoss是中间特征蒸馏(如MSE)。蒸馏训练时,学生模型同时计算任务损失和蒸馏损失。
13.4 压缩组合拳
我们最终方案:先剪枝(稀疏训练+剪枝+微调),再量化感知训练,最后用蒸馏进一步提升量化后的精度。mmrazor支持多阶段压缩,我们可以将算法串联。
14. 部署与运维:模型版本管理与监控
14.1 模型版本管理
随着超市商品更新,我们需要定期更新模型。我们使用DVC(数据版本控制)管理数据集,MLflow跟踪训练实验,每个模型都记录超参数、训练日志和评估指标。
14.2 边缘端监控
部署后,我们需要监控模型的推理性能(延迟、吞吐量)和预测质量。我们在Jetson上运行一个监控服务,定期采样摄像头图像,与人工标注对比,计算mAP和mIoU,发现异常时触发重新训练。
14.3 增量学习与在线适应
当新商品上架,我们需要快速更新模型而不遗忘旧知识。我们尝试了增量学习方法:在新数据上微调,但加上知识蒸馏损失(用旧模型作为教师)。mmrazor的蒸馏机制正好可以用于此场景,我们只需将旧模型作为教师,当前模型作为学生,在新数据上训练。
15. 总结:完整的技术闭环
通过这个超市项目,我们完整地走通了从数据、训练、压缩、部署到维护的全流程,并且每个环节都深入到了源码和工程细节。现在,当有人问你是否了解这些技术时,你可以自信地说:
“我不仅知道YOLO的原理,还能手写build_targets;不仅会用mmsegmentation,还能自定义数据集和损失;不仅会用mmdeploy导出模型,还能写重写规则解决算子兼容;不仅会用mmrazor剪枝量化,还能组合压缩策略提升部署效率。这些技术在我们超市系统中环环相扣,解决了真实场景的难题。”


