训练与新增类别教程

工程目录：D:\Edge\download\yolov5-v6.0\yolov5-v6.0 环境：Python 3.8 + YOLOv5 v6.0；依赖安装 pip install -r requirements.txt 部署目标：泰山派（RK3566）NPU

目录

1. 数据集长什么样
2. 怎么训练
3. 怎么评估
4. 新增识别目标要改哪些地方
5. 数据不够 / 类别不平衡怎么办
6. 部署链路：pt → onnx → rknn → 泰山派
7. 常见操作速查
8. 常见坑（实战踩坑）

一、数据集长什么样

dataset/
├─ train/{images,labels}    # 训练集
├─ valid/{images,labels}    # 验证集
├─ test/{images,labels}     # 测试集
├─ data.yaml                # 类别定义
└─ yolov5s.pt               # 预训练权重（可选放这里）

• 图片和标签同名一一对应（a.jpg ↔ a.txt）。
• 标签为 YOLO 格式，一行一个目标，数值全部归一化到 0~1：

<class_id> <x_center> <y_center> <width> <height>

• 没有目标的图片：标签文件留空（不要删文件）。
• data.yaml 决定类别：

path: D:/Edge/download/yolov5-v6.0/yolov5-v6.0/dataset
train: train/images
val: valid/images
test: test/images

nc: 6
names: ['person', 'hardhat', 'no-hardhat', 'safety-vest', 'no-safety-vest', 'fire']

names 的下标 = 标签里的 class_id，这个对应关系是全部工作的基准，别弄错。

二、怎么训练

1. 直接开跑（CPU 友好版）

python train.py --weights dataset/yolov5s.pt --cfg models/yolov5s.yaml --data dataset/data.yaml --epochs 50 --batch-size 8 --img 416 --device cpu --workers 4 --patience 15 --name my_v1

有独显就把 --device cpu 改成 --device 0，并把 --img 640 --batch-size 16 加回去。

2. 参数说明

3. ⚠️ CPU 训练时间（重点，别踩）

实测（约 400 张图、yolov5s、CPU）：

结论：CPU 训练一律用 --img 416 --epochs 50，别用 640/100 轮。 输入尺寸降到 416，计算量只有 640 的约 42%。

4. 训练结果在哪

runs/train/<name>/：

• weights/best.pt → 推理 / 部署用（验证集最优）
• weights/last.pt → 断点续训用（最后一轮）
• results.csv → 每轮的 loss / mAP 数值（可直接看）
• results.png / PR_curve.png / confusion_matrix.png → 评估图，训练跑完才生成

5. 续训

python train.py --resume runs/train/my_v1/weights/last.pt

三、怎么评估

1. 看每类指标（最重要）

python val.py --weights runs/train/my_v1/weights/best.pt --data dataset/data.yaml --img 416

输出示例：

Class     Images  Labels     P      R    mAP@.5  mAP@.5:.95
all         228    1161   0.742  0.672   0.705    0.399
person      228     155   0.416  0.194   0.188    0.079   ← 偏低
hardhat     228     372   0.819  0.704   0.762     0.41
fire        228      77   0.805  0.701   0.717    0.434

看 mAP@.5：

2. 看画面（可视化）

# 整个测试集
python detect.py --weights runs/train/my_v1/weights/best.pt --source dataset/test/images --img 416
# 单张图
python detect.py --weights runs/train/my_v1/weights/best.pt --source 你的图.jpg --img 416
# 摄像头实时
python detect.py --weights runs/train/my_v1/weights/best.pt --source 0 --img 416 --view-img

结果在 runs/detect/expN/。

3. 只看某个类

python val.py --weights runs/train/my_v1/weights/best.pt --data dataset/data.yaml --img 416 --classes 5

四、新增识别目标要改哪些地方 ⭐

核心原则：类别总数 nc 和类别顺序，在下面几处必须完全一致。

以「原来 5 类，加一个 fire → 6 类」为例：

1）改 data.yaml（必改）

nc: 6
names: ['person', 'hardhat', 'no-hardhat', 'safety-vest', 'no-safety-vest', 'fire']

names 的下标就是标签里的 class_id。

2）改标签 class_id（必改）

新类别编号必须是 0 ~ nc-1 内、且与 names 下标对应、不能和已有类别冲突。

例：火灾标签原来是 0/1/2，统一改成 5（= fire 在 names 中的下标）：

0 0.51 0.44 0.31 0.22   →   5 0.51 0.44 0.31 0.22

3）删缓存（改过标签就删）

删除 dataset/train/labels.cache、dataset/valid/labels.cache，否则 YOLOv5 读旧缓存，新标签不生效。

4）重新训练（必做）

nc 一变，网络检测头输出维度就变了，旧权重不能直接配新 data.yaml，必须重训。

5）部署端同步改（必改）

RKNN 转换脚本和板端后处理里的类别表要一致，否则类别名会错位：

CLASSES = ("person", "hardhat", "no-hardhat", "safety-vest", "no-safety-vest", "fire")

五、数据不够 / 类别不平衡怎么办

1. 先看差多少

用 val.py 看每类 mAP，再数各类样本数。经验值：

train : valid : test = 7 : 2 : 1（或 8:1:1），按每类 200 张 train 算，每类总量约 286 张。

2. mAP 低的常见原因

3. 去哪找数据

优先找已标注的数据集，直接合并：

• Roboflow Universe：https://universe.roboflow.com/search?q=PPE 搜 PPE detection、safety helmet、safety vest、no hard hat、fire detection 等； 下载时选 YOLOv5 / YOLOv8 (txt) 格式。
• Kaggle：搜同名关键词。

4. 合并多个数据集（关键）

不同数据集的类别顺序几乎肯定不一样，必须重映射 class_id 后再合并。步骤：

1. 建一张映射表（源类别 → 目标类别号），未列出的源类别丢弃；
2. 图片+标签拷入 dataset/train|valid|images|labels，同名成对；
3. 按上表改写标签的 class_id；
4. 删 labels.cache；
5. 重新训练。

两个必须处理的细节：

• 防数据泄漏：Roboflow 会导出同一张原图的增强副本（文件名前缀相同、只有 .rf.<hash> 不同）。切分 train/valid 时，同一原图的所有副本必须进同一个 split，否则验证集分数虚高。
• Windows 超长路径：某些来源的文件名超过 260 字符，Python 读取会报错。读写时用 \\?\ 前缀（"\\\\?\\" + os.path.abspath(p)），并另存为短文件名。

六、部署链路：pt → onnx → rknn → 泰山派

泰山派芯片 RK3566，只能用 .rknn 模型。转换必须在 PC（x86_64 Linux） 上做，板子上装不了转换工具。

1. 导出 ONNX（在 Windows 上）

先改两个文件（一次性）：

• models/yolo.py：Detect.forward 增加 RKNN 分支，输出 3 个原始张量：

def forward(self, x):
    z = []
    if os.getenv('RKNN_model_hack', '0') != '0':   # RKNN 导出路径
        for i in range(self.nl):
            x[i] = torch.sigmoid(self.m[i](x[i]))
        return x
    # ... 下面是原有逻辑，保持不变

• export.py：在 run() 里加一行

os.environ['RKNN_model_hack'] = 'npu_2'

然后导出：

python export.py --weights runs/train/my_v1/weights/best.pt --img 416 --batch 1 --include onnx --opset 12

产物 best.onnx，输出应为 3 个张量（以 416 为例）：

input:   [1, 3, 416, 416]
outputs: [1, 33, 52, 52]  [1, 33, 26, 26]  [1, 33, 13, 13]
# 33 = 3锚点 × (5 + nc=6)

2. 准备转换材料

需要 4 样东西（放同一个文件夹）：

3. Ubuntu 上转 RKNN

环境（conda 或 venv，Python 3.8）：

# 用现有环境即可，例如：
source ~/Val-Env/RK3566/bin/activate
python -c "import rknn; print(rknn.__file__)"

若没有环境：

python3.8 -m venv ~/rknn_env
source ~/rknn_env/bin/activate
pip install -U pip -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install rknn-toolkit2 -i https://pypi.tuna.tsinghua.edu.cn/simple

⚠️ 版本必须匹配：rknn-toolkit2 版本要 ≤ 板端 runtime 版本（板端用 strings librknnrt.so | grep version 查），否则板子上加载会报版本不匹配。板端是 1.5.x，就用 rknn-toolkit2==1.5.2。

转换脚本：

import os
from rknn.api import RKNN

ONNX, RKNN_OUT, DATASET = 'best.onnx', 'best.rknn', 'calib.txt'

rknn = RKNN(verbose=True)
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
            target_platform='rk3566',                  # 泰山派
            quantized_dtype='asymmetric_quantized-8')
assert rknn.load_onnx(model=ONNX) == 0, 'load_onnx failed'
assert rknn.build(do_quantization=True, dataset=DATASET) == 0, 'build failed'
assert rknn.export_rknn(RKNN_OUT) == 0, 'export failed'
rknn.release()
print('OK ->', RKNN_OUT, '%.1f MB' % (os.path.getsize(RKNN_OUT) / 1048576))

python rknn_convert.py      # 生成 best.rknn（YOLOv5s int8 约 4~8 MB）

4. 泰山派上运行

① 板子实况

② 查 runtime 版本（确认与 toolkit 匹配）

strings /usr/lib/librknnrt.so | grep -i "librknnrt version"

板端 ≥ 1.5.2 即可直接用 1.5.2 转出的模型（runtime 向下兼容旧模型）。

③ 传模型

# 在 Ubuntu(PC) 上执行；板子 IP 用 ip addr 查
scp ~/YOLO/my_all6/best.rknn huangwenbo@<板子IP>:~/rknn_yolov5/

④ 改 demo 的类别定义（在 rknn_yolov5_cam.cpp 里）

#define OBJ_CLASS_NUM 6                     /* 必须等于模型的 nc */
static const char *CLASSES[] = {
    "person", "hardhat", "no-hardhat",
    "safety-vest", "no-safety-vest", "fire"
};
const char *model_path = "./best.rknn";     /* 指向新模型 */

⑤ 后处理加类别过滤（只看火；模型不用改、不用重训）

/* 只保留 fire（class_id = 5），其余丢弃 */
if (class_id != 5) continue;

也可做成"火 + 安全帽"：

if (class_id != 1 && class_id != 5) continue;   /* 1=hardhat, 5=fire */

⑥ 编译并运行（沿用你原来的编译方式 / CMakeLists）

cd ~/rknn_yolov5
./rknn_yolov5_cam            # 或按原步骤重新编译后再跑

首次运行留意打印的 model input height/width，确认是 416（与导出一致）。

若换用 Rockchip 官方例程：rknpu2/examples/rknn_yolov5_demo， export LD_LIBRARY_PATH=./lib && ./rknn_yolov5_demo model/RK356X/best.rknn 图.jpg

七、常见操作速查

八、常见坑（实战踩坑）

训练类

• mAP 一直是 0：标签 class_id 超出 nc 范围，或 nc 与标签对不上。
• 新标签不生效：忘了删 labels.cache。
• 检测框类别名错位：names 顺序与标注时的类别顺序不一致。
• 图片被忽略：图片没有同名 .txt 标签，会被当成负样本或跳过。
• 改了 nc 却用旧权重：必须重新训练，不能只改 yaml。
• CPU 训练跑到天荒地老：用了 img 640 + epochs 100，务必改 img 416 + epochs 50。
• 验证分数虚高：Roboflow 增强副本被拆到了 train 和 valid 两边（数据泄漏），切分时要按原图分组。

部署类

• 在板子上装 rknn-toolkit2：装不了，转换只能在 PC(x86_64) 上做；板子只有 runtime。
• 板端加载报 version mismatch：toolkit 版本高于板端 runtime，降 toolkit 版本重转。
• python 找不到 rknn 包：环境没激活。用 find ~ -maxdepth 4 -name activate -path "*/bin/*" 找已建好的 venv，再 source .../bin/activate。
• RKNN 默认目标平台是 rk3588：必须显式写 target_platform='rk3566'，否则板端跑不了。
• ONNX 输出只有一个张量：忘了改 models/yolo.py 的 Detect.forward 或没在 export.py 里设 RKNN_model_hack。
• 量化后精度掉了：校准图太少或不具代表性，建议 100~300 张、覆盖各场景。
• 只想看火却一堆框：后处理没做类别过滤，加上 if (class_id != 5) continue;。

九、完整流程清单（照做即可）

阶段 0 · 环境准备

cd D:\Edge\download\yolov5-v6.0\yolov5-v6.0
pip install -r requirements.txt

阶段 1 · 准备数据集

1. 按 一、数据集结构 组织 dataset/
2. 确认 data.yaml 的 nc 与 names
3. 每类数量尽量均衡（train : valid : test = 7 : 2 : 1，每类 train ≥ 200）

阶段 2 · 训练

python train.py --weights dataset/yolov5s.pt --cfg models/yolov5s.yaml --data dataset/data.yaml --epochs 50 --batch-size 8 --img 416 --device cpu --workers 4 --patience 15 --name my_v1

阶段 3 · 评估

python val.py --weights runs/train/my_v1/weights/best.pt --data dataset/data.yaml --img 416
python detect.py --weights runs/train/my_v1/weights/best.pt --source dataset/test/images --img 416

mAP 达标（≥0.6）再往下走。

阶段 4 · 导出 ONNX（Windows）

1. 改 models/yolo.py 的 Detect.forward（加 RKNN 分支）
2. 改 export.py（加 os.environ['RKNN_model_hack'] = 'npu_2'）

python export.py --weights runs/train/my_v1/weights/best.pt --img 416 --batch 1 --include onnx --opset 12

阶段 5 · 转 RKNN（Ubuntu x86_64）

1. 新建工作目录，放入 best.onnx + calib/ + calib.txt + rknn_convert.py
2. 激活已有环境：source ~/Val-Env/RK3566/bin/activate（rknn-toolkit2 1.5.2）
3. 转换：

cd ~/YOLO/my_all6
python rknn_convert.py          # 产物 best.rknn，约 4~8 MB

阶段 6 · 板端部署（泰山派）

1. scp best.rknn huangwenbo@<板子IP>:~/rknn_yolov5/
2. 改 rknn_yolov5_cam.cpp：OBJ_CLASS_NUM / CLASSES[] / 模型路径
3. 后处理加类别过滤：if (class_id != 5) continue;
4. 编译运行，确认输入尺寸 416、框和类别名正确

十、本项目实际档案（可作参照）

经验结论：

• CPU 训练一律 img 416 + epochs 50，能省 3~5 倍时间；
• 数据均衡比堆量更重要：把每类拉到 200 张后，整体 mAP 从 0.403 → 0.705；
• 稀少/异构的类（如 person）光靠堆量不够，其瓶颈在标注口径不一致；
• 模型保持 6 类、在部署端做类别过滤，比"重训单类模型"更灵活；
• toolkit 版本必须 ≤ 板端 runtime，建议先在板子上 strings librknnrt.so | grep version 确认。