三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

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

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

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

我们目前的数据集包含:

  • 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 才会继续作为“疑似漏标”。

五、为什么还要分AUTOREVIEW

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

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

以 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

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

十一、GTAUTO的区分

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

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

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 优化里非常重要的正确性问题。

第三,把不需要继续留在GPUtensor尽快转到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使用情况。

← 返回列表