4. YOLOv5模型训练与RKNN转换实战

在工业视觉场景中，缺陷检测、目标定位和状态识别往往需要部署在边缘计算设备上。RK3588/RK3568等瑞芯微平台凭借其强大的NPU算力（RK3588达6 TOPS）成为工业网关和智能相机的理想选择。然而，从PyTorch训练的YOLOv5模型到能够在NPU上高效运行的RKNN格式，中间隔着数据标注、模型训练、ONNX导出、RKNN转换、精度验证五道关卡。本章将完整走通这条链路，重点解决“训练出的模型精度高但转换后掉点严重”“ONNX算子不兼容”“NPU量化后检测框偏移”等高频问题。

上图展示了从数据标注到RKNN验证的完整工程链路。实际项目中，每个环节都可能成为瓶颈：标注质量影响模型上限，训练超参决定收敛速度，ONNX导出决定算子兼容性，量化策略决定NPU上的最终精度。下面按阶段拆解，并给出可直接落地的操作命令和参数配置。

4.1 数据标注与预处理：工业场景的“脏数据”清洗

核心原理：YOLOv5使用归一化的txt格式标注，每行代表一个目标：class_id center_x center_y width height，坐标值均除以图片宽高。工业场景中，传送带上的工件、仪表盘读数、PCB缺陷等目标往往存在遮挡、反光、小目标密集等问题，标注质量直接影响后续所有环节。

操作流程：

1. 数据采集：使用工业相机（如海康MV-CA050-10GM）拍摄不同光照、角度、遮挡程度下的目标图片，建议至少2000张/类。
2. 标注工具：推荐X-AnyLabeling（支持自动标注辅助）或LabelImg。标注时遵循“框住完整目标，不包含背景干扰”原则。
3. 格式转换：若使用LabelImg的VOC格式，需转换为YOLO txt。可使用以下Python脚本批量处理：

import xml.etree.ElementTree as ET
import os

def voc_to_yolo(xml_path, out_dir, class_list):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    img_w = int(root.find('size/width').text)
    img_h = int(root.find('size/height').text)
    txt_name = os.path.splitext(os.path.basename(xml_path))[0] + '.txt'
    with open(os.path.join(out_dir, txt_name), 'w') as f:
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in class_list:
                continue
            cls_id = class_list.index(cls)
            bbox = obj.find('bndbox')
            x1 = float(bbox.find('xmin').text)
            y1 = float(bbox.find('ymin').text)
            x2 = float(bbox.find('xmax').text)
            y2 = float(bbox.find('ymax').text)
            cx = (x1 + x2) / 2 / img_w
            cy = (y1 + y2) / 2 / img_h
            w = (x2 - x1) / img_w
            h = (y2 - y1) / img_h
            f.write(f"{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")

案例拆解：某汽车零部件厂需要检测密封圈上的划痕，目标尺寸仅30×30像素，而图片为1920×1080。初始标注时，标注员将划痕周围的阴影也框入，导致模型误检率高达15%。分析判断：阴影与划痕在灰度上接近，但纹理不同。解决措施：① 在标注规范中明确“只框划痕本体，不包含阴影”；② 使用CLAHE（对比度受限自适应直方图均衡化）增强纹理；③ 将原图裁剪为640×640的滑窗进行训练，避免小目标被压缩丢失。

4.2 YOLOv5模型训练：迁移学习与超参调优

核心原理：工业场景通常数据量有限（几千张），直接从头训练YOLOv5容易过拟合。采用迁移学习策略：加载COCO预训练权重，冻结骨干网络前10层，只训练检测头和部分特征层。输入分辨率建议640×640，兼顾速度与精度。

关键训练配置（data.yaml）：

# data.yaml
train: ./datasets/train/images
val: ./datasets/val/images
nc: 3  # 类别数，例如：划痕、污渍、破损
names: ['scratch', 'stain', 'breakage']

训练命令：

python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --epochs 100 --device 0 --workers 8 --cache

超参调优要点：

案例拆解：某电子厂检测PCB板上的焊点缺陷，初始训练100轮后mAP@0.5为0.82，但验证集上小焊点（<20×20像素）漏检严重。分析：默认anchor针对COCO的通用目标，对微小目标不敏感。操作：在训练命令中加入 --autoanchor，重新聚类出适合小目标的anchor尺寸（如[4,5, 8,10, 12,15]）。同时将输入分辨率从640提升至960（--img 960），虽然推理速度下降约30%，但小目标mAP提升至0.91。

4.3 导出ONNX：算子兼容性与固定尺寸

核心原理：RKNN-Toolkit2对ONNX的算子支持有限，YOLOv5默认导出脚本中的torch.onnx.export可能产生NPU不支持的算子（如GridSample、部分Resize模式）。因此导出时需做以下处理：

1. 固定输入尺寸：RKNN转换时动态尺寸会显著增加转换难度和推理延迟，建议固定为640×640。
2. 简化模型：使用onnx-simplifier消除冗余的Shape/Reshape节点。
3. 修改检测头：YOLOv5的检测头包含nn.SiLU激活函数，RKNN支持SiLU，但需确认ONNX版本中导出为HardSwish或SiLU节点。

导出命令：

python export.py --weights runs/train/exp/weights/best.pt --img 640 --batch 1 --include onnx --opset 12 --simplify

验证ONNX输出：使用onnxruntime检查输出张量形状是否为[1, 25200, 85]（640×640输入，3个尺度，每个尺度3个anchor，共25200个候选框）。

4.4 RKNN-Toolkit2转换与量化：精度与速度的博弈

核心原理：RKNN-Toolkit2将ONNX模型转换为RKNN格式，并默认进行INT8量化。量化过程需要校准数据集（通常200-500张代表图片），用于统计每层激活值的动态范围。若校准集与真实场景分布差异大，会导致量化误差放大。

转换脚本（Python）：

from rknn.api import RKNN

rknn = RKNN()
# 配置量化参数
rknn.config(mean_values=[[0,0,0]], std_values=[[255,255,255]], target_platform='rk3588', quantized_dtype='w8a8')
# 加载ONNX
ret = rknn.load_onnx(model='yolov5s.onnx', input_size_list=[[1,3,640,640]])
if ret != 0:
    print("load onnx failed")
    exit(1)
# 构建模型（量化）
ret = rknn.build(do_quantization=True, dataset='./quant_data.txt')
if ret != 0:
    print("build failed")
    exit(1)
# 导出RKNN
ret = rknn.export_rknn('yolov5s.rknn')
rknn.release()

量化数据集准备：quant_data.txt中每行是图片路径，建议从验证集中随机抽取200张，覆盖不同光照、角度、目标大小。避免使用训练集（模型已见过，激活值分布过于理想）。

精度验证：转换后需在PC上模拟NPU推理，对比RKNN输出与PyTorch输出的mAP差异。

# 模拟推理
rknn.init_runtime(target=None)  # None表示PC模拟
outputs = rknn.inference(inputs=[img])
# 计算mAP，与PyTorch结果对比

案例拆解：某物流分拣线部署YOLOv5检测包裹上的标签，转换后mAP从0.95降至0.82，且小标签（30×30像素）几乎全部漏检。分析：量化校准集中包含大量大包裹图片，小标签的激活值被压缩到低比特范围。操作：① 重新构建校准集，确保小标签图片占比≥30%；② 使用混合量化，将检测头的前两层保持FP16，其余层INT8；③ 在rknn.config中设置quantized_algorithm='normal'并增加optimization_level=1。调整后mAP恢复至0.90，延迟仅增加2ms。

4.5 常见误区与总结复盘

误区一：标注数据不检查。训练前未可视化标注框，导致坐标偏移、类别错误等问题在训练后才暴露。建议训练前用python train.py --cache配合plot_labels工具生成标注分布图，检查每个类别的目标数量、尺寸分布。

误区二：ONNX导出后不验证。直接进行RKNN转换，遇到算子不支持时反复试错。正确做法是先用onnxruntime验证ONNX输出与PyTorch一致，再进入RKNN阶段。

误区三：量化校准集随意选。直接使用训练集前200张图片作为校准集，导致量化后精度虚高，部署到现场后掉点严重。校准集必须独立于训练集，且覆盖真实场景的分布。

误区四：忽略后处理差异。YOLOv5的NMS在PyTorch中实现，而RKNN模型只输出原始预测框。部署时需在C++/Python侧重新实现NMS，且注意RKNN输出的坐标是归一化的，需乘以输入尺寸还原。

总结复盘：本章完整走通了从标注到RKNN验证的链路。核心经验可归纳为：标注质量决定上限，量化策略决定下限。在工业项目中，建议预留20%的时间用于精度调优，而非全部投入训练。当RKNN转换后mAP下降超过2%时，优先检查校准集分布，其次考虑混合量化，最后才调整网络结构。