单类别模型自动补标工作流与迭代优化记录

一、我们现在要解决的核心问题

我们目前的数据集包含：

• person
• helmet
• vest
• tractor
• slipper
• smoking

共 6 个目标类别。

现有数据已经有人工标注，但随着数据量扩大，发现部分图片存在漏标问题。例如一张图片实际上有 4 个人，但原始标签只标了 3 个；或者图片中有安全帽、拖鞋、吸烟等小目标，但原始数据没有完整标出来。

如果全部重新人工检查，成本比较高。

所以我们的思路不是重新做一套伪标签，而是：

保留现有人工标签作为基准，再利用已经训练好的高质量单类别模型，对整个数据集进行一次“漏标扫描”，只补充模型高置信度发现、同时又没有和原有标签重合的目标。

目标是：原标签+模型辅助发现漏标+保守自动补标+人工抽样检查。

二、为什么使用6个单类别模型

目前我们分别有：

person model
helmet model
vest model
tractor model
slipper model
smoking model

不是同时把 6 个模型全部加载进显存，而是采用串行方式：

person 模型
↓
扫描 train / val / test
↓
释放模型

helmet 模型
↓
扫描 train / val / test
↓
释放模型

vest 模型
↓
扫描 train / val / test

tractor 模型
↓

slipper 模型
↓

smoking 模型

代码当前就是外层遍历类别，再在每个模型内部遍历 train、val、test。

假设整个数据集：

train + val + test = N 张

那么 6 个模型理论上就是：

6 × N

次 image-model inference。

例如有 20,000 张图：

20,000 × 6 = 120,000 次图片-模型推理

虽然计算量比较大，但是这样做有两个好处。

第一，不需要同时把 6 个模型放在 GPU 中，显存压力更小。

第二，每一个类别都可以使用自己的模型和自己的置信度阈值。例如 smoking、slipper 这种比较难的类别，可以和 person、helmet 使用不同标准。

当前默认阈值就是按类别分别设置的，例如：

person: auto=0.75 review=0.55
helmet: auto=0.75 review=0.55
vest: auto=0.75 review=0.55
tractor: auto=0.70 review=0.50
slipper: auto=0.65 review=0.45
smoking: auto=0.65 review=0.40

代码中已经将这些阈值独立配置。

三、安全原则：绝对不修改原标签

这个方案最核心的设计之一是：原始labels永远不动。

程序首先把：dataset/labels/复制一份成为：out_root/labels_autofill_v1/

代码就是先逐 split 复制原标签。

所以一开始：labels_autofill_v1=原始 labels

后面模型发现的新增标签全部只追加到：labels_autofill_v1/

代码的 append 操作也明确写的是输出标签目录，而不是原始目录。即使模型补错很多，也可以直接删除整个输出目录，原始数据完全不会受到影响。

四、模型到底怎么判断一个框是不是“漏标”

以 person 模型为例。

假设一张图实际有：

person A
person B
person C
person D

但是原标签只有：

person A
person B
person C

模型预测：

P1 → A，confidence=0.96
P2 → B，confidence=0.93
P3 → C，confidence=0.91
P4 → D，confidence=0.89

我们不能直接把这 4个框全部追加，否则 A、B、C 就会重复标注。

所以当前做了一层非常重要的：

Existing Label Matching

对于模型预测框：Prediction计算它和原来同类别GT框的最大 IoU。

例如：

P1 vs 原 GT 最大 IoU = 0.94
P2 vs 原 GT 最大 IoU = 0.89
P3 vs 原 GT 最大 IoU = 0.91
P4 vs 原 GT 最大 IoU = 0.02

默认：iou_existing = 0.50

如果：IoU >= 0.50

说明这个目标大概率原来已经标过：

matched_existing
→ continue
→ 不再补

当前代码就是这样处理。

所以：

P1 → 已标
P2 → 已标
P3 → 已标
P4 → 原标签附近没有 person

P4 才会继续作为“疑似漏标”。

五、为什么还要分AUTO和REVIEW

剩下的框并不是全部自动写进去，而是进一步按置信度分为：

低置信度
中置信度
高置信度

以 person 为例：

0 -------- 0.55 -------- 0.75 -------- 1
丢弃 REVIEW AUTO

也就是说：

低于review threshold

conf < 0.55模型自己都不太确定，直接忽略。

中等置信度

0.55 <= conf < 0.75

进入：REVIEW记录下来，但是不修改标签。

高置信度

conf >= 0.75进入：AUTO才允许自动追加到：labels_autofill_v1

当前代码中就是：

mode = auto if conf >= auto_thr else review

并且只有 batch_auto 最终会写入复制后的标签。

所以整个系统实际上是一个：

模型发现目标
│
▼
和已有 GT 比较
│
├── 已经标过 → 忽略
│
▼
真正疑似漏标
│
├── 低置信度 → 忽略
│
├── 中置信度 → REVIEW
│
└── 高置信度 → AUTO

这是一个比较保守的补标策略。

六、第二层IoU：解决“多个真实目标被错误去重”的问题

最早版本里还有一个问题：

用于判断“已有标签”的 IoU 阈值和“预测框重复”的 IoU 阈值使用得比较接近。

这会有风险。

例如两个人站得比较近：

person A box
person B box

可能：

IoU = 0.55

但是他们确实是两个真实的人。

如果把：

IoU > 0.5

直接当成重复预测，就可能误删一个人。

所以后面把两个概念彻底拆开：

existing IoU

--iou-existing = 0.50

回答的问题是：这个预测和原有人工标签是不是同一个目标？

candidate duplicate IoU

--iou-candidate-duplicate = 0.95

回答的问题是：两个新的模型预测是不是几乎完全相同的重复框？

当前代码已经把两者拆成两个参数。

并且只有新的 candidate 之间：IoU >= 0.95

才会认为是重复预测。这样可以避免因为两个物体靠得近，把真实目标错误删除。

七、IoU计算也做过专项排查

IoU 属于这套流程里非常关键的基础函数。

因为如果 IoU 算错，比如面积计算错误，就可能出现：

真实 IoU = 0.2
错误计算 IoU = 0.7

于是系统会误认为：这个目标原来已经标过。

最终造成：漏标依然没有补出来。

所以后续我们专门检查了 IoU 实现。

当前版本：

area_a = (ax2 - ax1) × (ay2 - ay1)
area_b = (bx2 - bx1) × (by2 - by1)

实现是正确的。

这类 Bug 比程序直接报错更危险，因为：程序可能完整运行结束，但是生成的数据质量是错的。所以现在我们的思路不仅是“代码能跑”，还要对：

IoU
class id
threshold
duplicate suppression
label append

这些关键数据逻辑做专项验证。

八、类别ID也不能直接写死

还有一个比较重要的优化是类别映射。

单类别模型内部可能都是：

class 0

比如：

person_best.pt:
0 = person

helmet_best.pt:
0 = helmet

slipper_best.pt:
0 = slipper

但是最终总数据集可能是：

0 person
1 helmet
2 vest
3 tractor
4 slipper
5 smoking

所以：helmet model 的 class 0不能直接写：0

否则最终会被训练程序解释成：person

因此现在程序从：data.yaml动态读取真实类别顺序和 class ID，并检查 nc、names 是否匹配。最终 candidate 创建时使用的是：

dataset_class_ids[class_name]

而不是单模型自己的类别编号。

这个优化主要是为了防止静默的数据污染。

九、auto_samples为什么后来重新设计

最初可视化逻辑是：

一个 Candidate
→ 读取一次图片
→ 只画这个 Candidate 的一个框
→ 保存

于是出现了一个明显问题。

例如：一张图片有 4 个 person但是一个自动补标 candidate 对应一张输出图，所以最终打开：auto_samples/person/xxx.jpg可能只看到一个框。

这会给人工检查造成误导：看起来像模型只检测到了一个人。

还有一个问题：如果两个 candidate 来自同一张图，并且：

conf = 0.92341
conf = 0.92336

旧文件名只保留 3 位：0.923就可能发生文件覆盖。

所以后面把抽样单位从：Candidate / box改成：unique image

CandidateWriter 现在按：mode + class + image进行分组，只保存高置信度的有限数量不同图片。因此：

--draw-auto-samples 80

表示的是：每个类别最多抽 80 张不同图片。

而不是：每个类别只保存 80 个框。

十、为什么仍然按类别分目录抽样

我们的目的不是随机看所有模型结果，而是分别评估：

person 补得怎么样？
helmet 补得怎么样？
vest 补得怎么样？
tractor 补得怎么样？
slipper 补得怎么样？
smoking 补得怎么样？

所以仍然保持：

auto_samples/
├── person/
├── helmet/
├── vest/
├── tractor/
├── slipper/
└── smoking/

当前程序确实按类别分别维护 sample group，并输出到对应类别目录。

这里需要说明一个当前版本的实现细节：

目前目录是按类别抽样的，但draw_sample_group()会先画这张图片完整的merged label状态，再高亮当前候选。

也就是说 person/ 中抽到的图片一定是因为存在 person 补标候选，但画面中仍可能看到其他类别标签。当前代码在绘制 merged boxes 时没有按 class_id 再过滤。

如果我们希望以后做得更纯粹，可以进一步改成：

auto_samples/person/
只显示 person 的 GT + person 的 AUTO

auto_samples/helmet/
只显示 helmet 的 GT + helmet 的 AUTO

这个属于下一步比较小的可视化优化，不影响真实补标结果。

十一、GT和AUTO的区分

可视化并不是简单把所有框都画一样。

当前会根据标签来源标记：

GT person
AUTO person

原有标签：

GT

自动新增标签：

AUTO

当前代码还使用不同线宽：

GT → thickness 2
AUTO → thickness 3

然后当前触发抽样的 candidate 还会额外显示：

AUTO person 0.92

这部分逻辑已经在可视化函数中实现。

后续还可以继续优化成：

GT → 一种颜色
AUTO → 另一种颜色

让老师或标注人员一眼就能看出来哪些是人工标签、哪些是模型新增标签。

十二、为什么经常出现显存OOM

目前机器显存已经比较大，但仍然会出现 CUDA OOM。

这个问题并不能简单理解成：显存容量小。

因为推理阶段显存峰值与多个因素有关：

输入尺寸 imgsz
batch size
模型大小
中间 feature map
YOLO NMS 前候选数量
PyTorch caching allocator
GPU 上还存活的 tensor
多模型切换后的缓存和碎片

我们当前：

imgsz = 832
batch = 32

本身就会产生比较大的推理峰值。当前默认配置可见于参数定义。

十三、当前已经实现的显存优化

1.一个时间只加载一个模型

不是：

person
helmet
vest
tractor
slipper
smoking

全部同时进入 GPU

而是：

load person
→ scan
→ delete person

load helmet
→ scan
→ delete helmet

每个类别完成后执行：

del model
gc.collect()
torch.cuda.empty_cache()

当前代码已经这样做。

这是目前最重要的显存控制策略。

2.使用FP16

如果设备是 CUDA：

half=True

使用 FP16 推理。

当前代码专门限制：

只有 CUDA 才启用 half precision。

避免 CPU/MPS 等设备误传 FP16。

FP16 对 YOLO 推理通常可以明显降低 activation 和 tensor 的显存占用。

3. batch推理

不是逐张：

image1
image2
image3
...

而是：

32 张
→ 一次 batch inference

这样 GPU 利用率更高。

4.使用stream=True

当前 model.predict()：

stream=True

也就是结果逐步消费，而不是要求所有结果长期堆在 Python 端。

5.每批和模型切换后进行对象清理

当前版本在 batch 结束后：

del results
gc.collect()
torch.cuda.empty_cache()

模型完成后再次执行清理。

这是目前采取的比较积极的显存回收方案。

十四、CPU内存也做了控制

不仅需要考虑 GPU 显存，还需要考虑普通 RAM。

如果一次扫描几十万图片，而把：

所有预测
所有 Candidate
所有可视化信息

全部保存在 list 中，内存会不断增长。

所以现在 CandidateWriter 的思路是：

Candidate 产生
↓
立即写 CSV

而不是：

几百万 Candidate 全部保存在 RAM
↓
最后再写 CSV

当前三个 CSV：

candidates_auto.csv
candidates_review.csv
candidates_all.csv

都是打开文件后持续流式写。

内存中只保留：用于最终可视化的有限数量 top sample。

例如：

auto 每类最多 80 张
review 每类最多 500 张

因此候选数量即使变成几十万，RAM 也不会和候选数量线性增长。

十五、标签Cache也按单类别释放

当前不是一次把：

6 个类别所有 GT

全部加载到 RAM。

而是在 person 模型开始时：

只 preload person labels

person 完成后：

del label_cache

helmet 再重新加载 helmet labels。

当前 preload_labels() 也是只筛选当前 class_id。

这个思路同样是：

计算时间稍微增加，但降低长期内存占用。

十六、磁盘空间也做了优化

如果最后需要形成一个标准 YOLO 可训练数据集，可以使用：

--materialize-dataset

程序生成：

trainable_dataset/
├── images/
├── labels/
└── data.yaml

其中图片默认不是再完整复制一份，而优先：

hardlink

这样：

原始图片
+
新训练数据集图片

在文件系统层面可以共享同一份数据块，大幅减少磁盘占用。

如果 hardlink 不支持，再自动 fallback 到 copy。当前逻辑已经实现。

十七、当前OOM还有什么不足

这个地方汇报时建议如实说。

我们之前已经讨论过一个更好的方案：

batch=32
OOM
↓
自动 batch=16
↓
还 OOM
batch=8
↓
直到成功

但是我重新检查当前实际脚本后发现：

这个“自动降batch并重试”的机制目前还没有真正落进当前可下载版本。

当前代码仍然是：

发生 CUDA OOM
↓
empty_cache
↓
抛出异常
↓
提示用户把 batch 减半重新运行

当前源码也明确写的是：

Retry with --batch ...

所以汇报时应该说：

目前已经有单模型串行、FP16、stream 推理和主动显存清理；下一轮准备进一步实现 adaptive batch，在发生 OOM 时让 batch 自动从 32 降到 16、8、4、2、1，而不是整个任务退出。

十八、下一轮显存优化方案

下一版我准备重点优化以下几个地方。

第一，自适应batch

把：--batch 32理解成：最大允许 batch。运行：

32
↓ OOM
16
↓ OOM
8
↓ success

然后从失败位置继续运行。

这样不同类别可以自动找到适合自己的 batch。

例如：

person → 32
helmet → 32
vest → 16
tractor → 32
slipper → 16
smoking → 8

因为不同模型本身的显存需求可能不一样。

第二，OOM重试必须做成事务式

这里还有一个比较容易忽略的问题。

使用：

stream=True

以后可能：

batch 32
前 15 张已经返回结果
第 16 张 OOM

如果前 15 张已经写进：

CSV
labels

然后 batch 缩小重新跑，就会产生重复标注。

所以正确设计应该是：

整个 batch 推理完成
↓
确认没有 OOM
↓
统一 commit CSV / labels

如果中途 OOM：

整个 batch 不提交任何结果
↓
降低 batch
↓
重新运行

也就是类似数据库：

batch inference transaction

这是下一轮 OOM 优化里非常重要的正确性问题。

第三，把不需要继续留在GPU的tensor尽快转到CPU

目前：

xywhn
confidence
class
IoU

部分后处理仍然在 GPU tensor 上进行。

实际上模型 forward 和 NMS 完成后，这些数据量已经非常小。

后续可以：

GPU prediction
↓
NMS
↓
tensor.cpu()
↓
CPU 上完成 IoU / threshold / candidate 构建

这样 GPU 只负责最需要 GPU 的模型 forward。

可以减少 tensor 生命周期重叠造成的峰值显存。

第四，不再每个batch都强制empty_cache

当前实现比较激进：

每个 batch
↓
torch.cuda.empty_cache()

这可以帮助释放 cache，但频繁调用也会增加 CUDA allocator 重新申请显存的开销。

后续可以优化成：

普通成功 batch
→ 不 empty_cache

OOM
→ empty_cache

模型切换
→ empty_cache

让 PyTorch 正常利用 caching allocator。

目标是在：

稳定显存

和：

推理吞吐

之间取得平衡。

第五，监控allocated / reserved / free

以后发生 OOM 时不能只输出：

CUDA out of memory

而应该记录：

GPU free memory
PyTorch allocated
PyTorch reserved
peak allocated
batch
imgsz
model
class
split

这样可以判断：

是真正 batch 太大
还是其他程序占 GPU
还是 reserved 过高
还是长任务出现碎片

这样显存优化就从：

凭经验调 batch

变成：

有监控数据地优化。

十九、CPU内存的下一步优化

目前最大的 CPU cache 是当前类别的：

label_cache

如果以后数据量从几万张扩大到几十万甚至百万张，仍然可以继续优化。

现在是：

一个 class
→ train + val + test label 全部 cache

未来可以改成：

person/train
→ 只 cache train
→ 完成释放

person/val
→ cache val
→ 完成释放

或者：

按图片 lazy load label

这样进一步降低 RAM。

不过当前数据规模下，优先级没有 GPU OOM 高。

二十、计算效率上的一个天然问题：同一张图片要推理6次

当前方法最大的问题不是内存，而是计算量。

假设一张图：

0001.jpg

最终需要经过：

person model
helmet model
vest model
tractor model
slipper model
smoking model

也就是 decode / preprocess / inference 六次。

当前代码虽然把图片路径只扫描一次并复用，避免了重复遍历目录，

但真正的神经网络 inference 仍然是六遍。

这是单类模型方案天然的代价。

我们暂时接受这个代价，因为现阶段更关注：

补标准确率
类别可控性
不同类别模型独立优化

而不是单纯追求最快速度。

如果后续数据量变得非常大，可以考虑：

6 个单类 teacher
↓
产生高质量补标数据
↓
训练一个统一的 6-class student model

以后新的数据直接：

1 个模型扫描一次

就能完成 6 类检测。

这样相当于把当前单类别模型当成：

teacher ensemble /专项教师模型

再通过补标后的数据反哺统一模型。

二十一、最终我们会得到哪些结果

程序运行完成后，输出目录主要包括：

out_root/
│
├── labels_autofill_v1/
│
├── candidates_auto.csv
├── candidates_review.csv
├── candidates_all.csv
│
├── auto_samples/
│ ├── person/
│ ├── helmet/
│ ├── vest/
│ ├── tractor/
│ ├── slipper/
│ └── smoking/
│
├── review_images/
│
├── summary.json
├── summary.txt
│
└── trainable_dataset/ # 可选

其中：

labels_autofill_v1

最重要。

原始人工标签
+
高置信度自动补标签

但是原始 dataset/labels 完全没有修改。

candidates_auto.csv

记录：

到底自动补了哪些框。

包括：

image
class
confidence
IoU
bounding box

candidates_review.csv

记录：

哪些候选模型认为可能漏标，但是置信度不足以自动修改。

方便人工二次审核。

auto_samples

用于抽样检查：

每个类别自动补标的实际效果。

summary

统计：

每类扫描多少图片
产生多少 prediction
多少与原标签匹配
多少 duplicate
最终 auto 多少
review 多少
error 多少

当前 summary 已经记录这些核心指标。

二十三、这样总结整个迭代过程

最初我们只是希望：用单类模型扫描数据，把漏标补回来。

但真正做以后发现它并不是简单的：

model.predict()
→ 写 txt

中间逐步暴露出很多工程问题。

第一阶段解决的是数据安全问题：

不修改原始标签，而是复制一套 labels，在复制本上补。

第二阶段解决的是重复标注问题：

通过同类别 IoU 判断预测是不是原来已经标过。

第三阶段解决的是错误自动补标问题：

使用 auto/review 双阈值，只允许高置信度自动写入，中间候选人工审核。

第四阶段解决的是真实目标被误去重的问题：

existing IoU 和 candidate duplicate IoU 分离，后者提高到 0.95。

第五阶段解决的是类别ID污染问题：

单类模型 class 0 不能直接写回总数据集，而是动态读取 data.yaml 的真实 class id。

第六阶段解决的是可视化误导问题：

从“一个框一张图”改成“按类别、按唯一图片抽样”，避免多目标图片只显示一个框和文件覆盖。

第七阶段解决的是运行安全问题：

对 output path 做保护，防止 --force 配错路径把原数据或者模型删掉。当前代码已经包含这类安全校验。

第八阶段开始重点解决大规模推理的GPU / RAM问题：

已经采用：

单模型串行
FP16
batch inference
stream prediction
CSV streaming
bounded visualization samples
模型切换释放
CUDA cache 清理

下一轮重点实现：

adaptive batch
OOM transaction retry
GPU tensor 及时搬 CPU
显存状态监控
减少过度 empty_cache

所以现在这项工作已经从：一个简单的自动标注脚本 逐渐变成：

一个面向大规模YOLO数据集的、安全可回滚、分级补标、可审计、可视化质检，并考虑GPU/CPU资源控制的数据增强流水线。

二十四、老师如果问“为什么不直接让模型全部重标？”

我会回答：我们没有选择全量覆盖原标签，因为模型预测不可能比所有人工标签都可靠。

现有人工标注仍然作为：

Ground Truth baseline

模型只负责：

发现人工可能漏掉的目标。

所以是：

人工 GT 为主
+
模型补漏为辅

而不是：

模型重新定义整个数据集。

这样风险小得多。

二十五、老师如果问“为什么不一次加载6个模型？”

可以回答：一次加载 6 个模型虽然可能减少模型切换次数，但是：

模型权重+CUDA context+batch activation+NMS tensor会共同占用显存。

尤其我们输入尺寸是 832，并且希望保持较大的 batch，因此目前选择：

时间换空间。

一次只保留一个模型，整个模型完成后释放，再加载下一个。

这样对长时间大数据任务更稳定。

二十六、老师如果问“40GB显存为什么还会OOM？”

可以回答：显存并不是只存模型权重。实际 peak memory 主要可能来自：

model weights+batch × 832 × 832 输入+网络多层 feature maps+prediction tensors+NMS intermediate tensors+PyTorch reserved memory+之前仍未释放的 GPU tensor

因此即使总显存很大：

batch=32+imgsz=832

仍可能在某些图片目标特别多、模型更复杂或者 CUDA 显存碎片较严重时出现峰值 OOM。

所以我们的方向不是简单换更大的显卡，而是：

让推理流水线具有自适应资源管理能力。

最终目标应该是：

用户只指定 max_batch=32

程序自己根据显存：
32 → 16 → 8 → ...

找到当前模型和数据能够稳定运行的最大 batch。

二十七、最后一句话总结

这项工作的核心不是“让模型替代人工标注”，而是：

利用多个高质量单类别模型作为teacher，对已有人工标注数据进行保守的漏标发现和增量修复；通过原标签保护、IoU匹配、AUTO/REVIEW分级、候选去重、类别映射、按类别可视化抽查以及资源控制，使补标结果既能提升数据完整性，又能够回滚、检查和追踪。

下一步重点是继续把显存管理从：人工遇到 OOM 后调 batch升级到：

系统自动感知OOM、自适应batch、事务式重试，并记录GPU使用情况。