1. 整体故事：我要开一家“未来超市”

假设我们团队接了一个项目：为一家连锁超市打造AI视觉系统，实现无人结算、货架智能管理、顾客行为分析。硬件是Jetson AGX Orin（性能比Nano强，但依然资源受限），要求所有模型在边缘实时运行（摄像头30fps）。系统包括：

  • 目标检测：识别顾客拿取的商品，并关联到购物车。
  • 语义分割：分析货架商品覆盖情况、顾客手部与货架的交互区域。
  • 模型部署：所有模型必须转换成能在边缘高效运行的格式，并做极致优化。
  • 模型压缩：为了达到实时性，需要对模型进行剪枝、量化，甚至知识蒸馏。

这个项目涉及的技术栈正好是你提到的所有内容。下面我们一步步深入。

2. 目标检测：YOLO源码深度剖析

2.1 为什么选择YOLOv5（或者v8）？

YOLO系列是目前工业界最常用的实时检测器。我们选择YOLOv5，因为它平衡了速度和精度，且工程化做得很好（代码易读、部署友好）。但作为工程师，不能只会调包，必须理解源码里的每个模块。

2.2 YOLOv5 网络结构细节

YOLOv5的网络可以分为：

  • Backbone：CSPDarknet53，但加入了Focus层和CSPNet结构。
    • Focus层：把输入图每隔一个像素取一个值，得到4张下采样2倍的特征图，再拼接。目的是减少计算量同时保留信息。但Focus层在部署时容易出问题，很多推理引擎不支持，可以用一个步长为2的6x6卷积替代。
    • CSPNet：Cross Stage Partial，将特征分成两部分，一部分经过卷积，另一部分直接concat，减少计算量，增强梯度。代码里通过BottleneckCSP或C3模块实现，内部有残差连接。
  • Neck：PANet（Path Aggregation Network），包括自上而下的FPN和自下而上的PAN，融合不同尺度特征。源码中通过YOLOv5Neck类实现，主要操作是上采样、卷积、concat。
  • Head：三个输出头，分别对应大中小物体，每个输出头输出(batch, anchors, grid_h, grid_w, (5+num_classes))，其中5是边界框的4个坐标（相对grid cell的偏移）和1个物体置信度。

2.3 训练源码核心：损失函数与正负样本分配

YOLOv5的损失计算是源码最复杂的部分，包含三部分：分类损失、置信度损失、框回归损失。

框回归损失：早期YOLO用MSE，后来用IoU系列。YOLOv5用的是CIoU Loss，它考虑了：

  • IoU（重叠面积）
  • 中心点距离（归一化到对角线长度）
  • 长宽比一致性

公式：

1
2
3
4
CIoU = IoU - (distance_center^2 / c^2) - α * v
v = (4/π^2) * (arctan(w_gt/h_gt) - arctan(w/h))^2
α = v / (1 - IoU + v)
Loss = 1 - CIoU

源码在utils/loss.py中，可以看到具体的张量计算。

正负样本分配：YOLOv5使用了一种称为跨网格匹配的策略。不同于Faster R-CNN的固定IoU阈值，它：

  1. 对每个GT，计算它与所有anchor的宽高比（w_ratio, h_ratio），如果比例在[1/4, 4]之间，就认为这个anchor能匹配该GT。
  2. 然后，根据GT的中心点，选择最近的几个grid cell（最多3个）也作为正样本，增加正样本数量。
  3. 每个正样本需要预测该GT的类别、框偏移和置信度（目标为1）。
  4. 负样本是所有未被选中的grid cell，只计算置信度损失（目标为0）。

这种分配方式在源码中通过build_targets函数实现，它利用广播和索引操作高效生成正样本mask。

数据增强：YOLOv5的Mosaic增强很经典，将4张图拼成一张，丰富了小物体和背景。但Mosaic在训练后期可能影响精度，所以会在最后epoch关闭。

2.4 推理源码细节

推理时，YOLOv5的输出需要解码：

  • 边界框中心相对于grid cell的偏移，通过sigmoid限制在0~1之间，再加上grid坐标。
  • 宽高通过exp(预测值)乘以anchor宽高。
  • 最后用NMS（非极大抑制）去除冗余框。

源码中，Detections类处理这些后处理，但为了部署，我们通常会把这些操作固定在模型内部（如加入sigmoid和decode层），或者用ncnn的层来实现。

3. 语义分割：从FCN到UNet再到mmsegmentation

3.1 为什么用UNet？

在超市场景中，我们需要对货架进行像素级分析，比如区分“商品”、“货架板”、“空隙”。这要求模型能保留高分辨率细节，UNet的U形结构和跳跃连接正好擅长这一点。

3.2 FCN与UNet的对比

  • FCN：将全连接层替换为卷积，通过转置卷积上采样，但上采样的结果比较粗糙，因为丢失了细节。它提出了跳跃连接的想法，但只是简单相加。
  • UNet：更彻底地使用跳跃连接，将encoder每个阶段的特征与decoder对应上采样的特征拼接，这样decoder能直接利用高分辨率信息，恢复精细边缘。

在mmsegmentation中，UNet的实现是模块化的：backbone可以是ResNet，decoder是UNetDecoder。源码里，UNet类继承自BaseSegmentor，它有一个_decode_head和_auxiliary_head。解码器的具体实现在mmseg/models/decode_heads/unet_head.py，可以看到如何逐层上采样并拼接。

3.3 源码深度：mmsegmentation的模块化设计

mmsegmentation是OpenMMLab体系的一部分，它的设计哲学是可配置、可扩展。理解它的源码能让你快速定制模型。

  • 注册器机制：所有模型组件（backbone、head、loss等）都通过Registry管理。比如@BACKBONES.register_module()装饰器，将类注册到BACKBONES字典。配置文件里通过type='ResNet'就能实例化。
  • 配置文件系统：使用mmcv.Config，支持从python文件加载配置，并且可以覆盖参数。比如想修改UNet的head类型，只需改配置文件中的decode_head字段。
  • 训练流程：Runner控制训练循环，调用model.train_step。每个模型需要实现forward_train和loss方法。
  • 数据流水线：通过Pipeline实现数据增强，比如LoadImageFromFile、Resize、RandomFlip等，可以自由组合。

理解这些，你就可以轻松地在mmsegmentation中添加新的backbone，比如把YOLOv5的CSPNet作为分割backbone，然后写一个适配器。

3.4 UNet在超市例子中的挑战与优化

  • 类别不平衡：货架空隙可能远多于商品，需要调整损失函数权重。mmsegmentation支持class_weight参数，可以在CrossEntropyLoss中设置。
  • 边缘模糊：可以增加边界感知损失，比如让模型在边缘像素上加重惩罚。这需要自定义loss，mmsegmentation允许你注册新的loss，并在配置里启用。

4. 模型部署：mmdeploy + ONNX + ncnn 的工程实践

4.1 mmdeploy 的作用与源码

mmdeploy把OpenMMLab的模型（如mmdetection、mmsegmentation）一键导出为各种后端格式。它的核心思想是重写PyTorch算子，使其兼容ONNX或其他后端。

源码结构：

  • mmdeploy/backend/：封装不同后端（onnxruntime、ncnn、tensorrt）。
  • mmdeploy/rewriter/：定义了各种算子的重写规则。比如adaptive_avg_pool2d在ONNX中可能被转换为GlobalAveragePool，但某些版本不支持动态输入，需要重写。
  • mmdeploy/codebase/：针对不同codebase（如mmdet）的导出适配器，负责提取模型输出、处理后处理等。

例如，导出YOLOv5时，mmdeploy会自动将模型的Detect层拆分成多个输出，并把后处理（decode + nms）移到外面，或者用后端支持的nms替代。

4.2 ONNX 导出细节

ONNX是中间表示，但PyTorch算子并非全部能直接转成ONNX。常见坑：

  • 动态尺寸：输入尺寸可能变化，导出时需设置dynamic_axes，比如dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}}。
  • 控制流：PyTorch中的if语句可能导不出，需保证模型计算图是静态的。
  • 算子支持：有些算子（如torch.where）在旧版ONNX中不支持，需要改写或用onnx-simplifier优化。

mmdeploy通过重写规则自动处理这些。比如YOLOv5的Focus层，它会被重写为几个切片和concat操作，最终转换为ONNX支持的Slice+Concat。

4.3 ncnn 部署优化

ncnn是腾讯的移动端推理框架，它针对ARM架构做了深度优化。部署时，我们先把ONNX转成ncnn格式：

1
onnx2ncnn model.onnx model.param model.bin

但转换后可能遇到算子不支持或效率低的问题。比如：

  • 上采样：UNet的双线性插值在ncnn中用Interp层，但需要指定resize_type=2。
  • BN层融合：ncnn的optimize工具可以融合Conv+BN，减少计算量。
  • FP16存储：ncnn支持加载FP16的bin文件，可减少内存。

在源码层面，ncnn的层都继承自Layer，实现了forward函数。如果需要自定义算子，可以仿照现有实现写一个，然后注册。

实际部署经验：我们曾遇到YOLOv5的SiLU激活函数在ncnn中较慢，将其替换为ReLU或LeakyReLU并微调，速度提升20%。

4.4 边缘部署的性能调优

在Jetson AGX Orin上，我们还可以用TensorRT，它比ncnn更快。但这里只讨论ncnn。调优手段：

  • 多线程：ncnn支持openmp，设置opt.num_threads。
  • 内存复用：通过ncnn::Extractor的set_memory_allocator复用内存。
  • 量化：ncnn支持INT8量化，需要校准数据集。

5. 模型压缩：mmrazor的量化、剪枝、蒸馏

5.1 mmrazor 框架概述

mmrazor是一个可扩展的模型压缩工具箱，支持剪枝、量化、蒸馏。它的设计类似于mmsegmentation，也是模块化+注册器。核心概念：

  • 算法（Algorithm）：定义压缩流程，如剪枝算法。
  • 组件（Components）：如剪枝器（Pruner）、量化器（Quantizer）、蒸馏器（Distiller）。
  • 可重写模块（Mutable）：定义模型中哪些部分可变，比如哪些通道可剪。

5.2 剪枝：基于BN层gamma的通道剪枝

常用方法是利用BN层的缩放因子gamma作为通道重要性指标。训练时对gamma施加L1正则化，让不重要的通道gamma趋近0，然后剪掉。

源码细节：

  • 在mmrazor中，剪枝器Pruner继承自BaseAlgorithm。它会在训练过程中，通过hooks在每次迭代后更新mask，或者使用LearnablePruner让mask可学习。
  • 剪枝时，需要生成新模型结构。mmrazor提供了MutableChannelUnit来标记每个卷积层的输出通道是否可剪。然后通过Pruner.get_pruned_model生成更窄的模型。
  • 微调阶段，把剪枝后的模型重参数化，即把原模型权重映射到新模型。注意跳跃连接的对齐。

在超市项目：我们想剪枝YOLOv5的backbone。需要定义哪些层可剪（比如所有Conv后面的BN），然后训练带L1正则的模型。mmrazor可以自动插入L1损失。

5.3 量化：量化感知训练（QAT）

量化把FP32模型转为INT8，加速4倍，内存减4倍。但后训练量化（PTQ）可能掉点严重，所以用QAT。

原理：在模型中插入伪量化节点（FakeQuant），模拟量化误差。这些节点记录输入的最大最小值，计算出scale和zero point，并在前向时对数据进行量化再反量化，保持梯度可导。

mmrazor实现：

  • Quantizer类负责替换网络中的Conv、Linear等为量化版本（比如nn.Conv2d换成mmrazor.ops.QuantConv2d）。
  • 量化器维护量化参数（scale, zero point），可以通过EMA（指数移动平均）统计动态范围。
  • 训练时，伪量化节点会截断梯度，但整体可以正常反向传播。

源码注意：需要处理好对称量化vs非对称量化、per-tensor vs per-channel。在ncnn中，INT8推理通常用对称量化（scale为正），且每个卷积核可能用不同的scale（per-channel）。

5.4 知识蒸馏

为了进一步提升小模型精度，可以用大模型（如YOLOv5x）蒸馏小模型（YOLOv5s）。蒸馏损失包括：学生预测与教师预测的KL散度（针对分类）、特征蒸馏（让中间特征对齐）。

mmrazor提供了蒸馏算法，如ConfigurableDistiller，可以配置哪些层的输出用于蒸馏。在超市场景，我们可以先训练一个大的YOLOv5l，然后蒸馏到剪枝后的YOLOv5s。

6. 串联整个系统：超市里的完整技术栈

让我们把以上技术串成一个完整的落地流程：

  1. 数据采集与标注：从超市多个摄像头采集图像，标注商品边界框（检测）和像素级标签（分割）。使用labelImg和labelme，转为COCO格式。
  2. 训练检测模型：用YOLOv5源码训练，重点调参（anchor、学习率、数据增强）。理解源码中的loss.py，调整正负样本分配策略（比如放宽宽高比阈值，适应细长商品）。
  3. 训练分割模型：用mmsegmentation训练UNet，利用其模块化设计，尝试不同的backbone（如Swin Transformer），并自定义边界损失。
  4. 模型导出：使用mmdeploy将两个模型导出为ONNX，解决YOLO的Focus层和UNet的上采样问题。然后转为ncnn格式。
  5. 性能分析：在Jetson上跑ncnn模型，发现速度只有10fps，需要压缩。
  6. 模型压缩：
    • 先对YOLOv5进行剪枝：用mmrazor的BNScalePruner，训练时加入稀疏正则，剪掉50%通道，微调恢复精度。
    • 再对UNet进行量化感知训练：用mmrazor的QATQuantizer，转为INT8，精度损失<1%。
    • 最后对YOLO进行蒸馏：用大模型指导小模型，提升剪枝后的精度。
  7. 重新部署：把压缩后的ncnn模型部署，速度达到25fps，接近实时。再用ncnn的Vulkan加速，达到30fps。
  8. 持续迭代：收集边缘案例，用增量学习更新模型。mmrazor支持在线蒸馏，可以在新数据上微调而不遗忘旧知识。

7. 深度答疑

最后，总结几个能体现深度的点，这些都是在实际工程中踩过的坑或深入理解的细节：

  • YOLOv5的正负样本分配源码：在utils/loss.py的build_targets中，你会看到如何用torch.meshgrid生成网格，然后用torch.where和torch.logical_and高效分配样本。理解后可以自定义分配策略，比如针对小物体增加匹配数量。
  • UNet的跳跃连接实现：在mmsegmentation的unet_head.py中，_forward方法里，每一步上采样后都要和对应的encoder输出拼接。注意拼接前要对encoder输出做卷积（reduce_channels）对齐通道数。
  • ONNX导出时的动态轴处理：在mmdeploy的onnx_helper.py中，有函数add_dynamic_axes，它会递归遍历计算图，为所有输入输出设置动态轴。但要注意有些算子（如Resize）不支持动态scale，需要重写。
  • ncnn中的卷积优化：ncnn的卷积实现是高度优化的，比如convolution_arm.h中，针对不同内核大小（3x3, 1x1）有专门汇编。你可以看到Winograd、FP16等优化。如果自己写自定义层，需要参考这些代码。
  • mmrazor的剪枝后模型重参：剪枝后，原模型的权重如何映射到新模型？mmrazor通过ModelConnector类，根据剪枝mask，从原模型state_dict中提取对应通道的权重，赋值给新模型。这涉及复杂的索引操作，比如torch.index_select和维度置换。
  • 量化中的calibration：PTQ时需要校准数据集来统计激活范围。mmrazor的calibrate方法会跑几次前向，收集各层的min/max，然后平滑处理。源码中会看到Observer类（如MinMaxObserver）负责收集统计量。