✅ YOLO获取COCO指标(4): 训练(Train)启用COCO API评估(实时监控AP指标)| 发论文必看! | Ultralytics | 小白友好
文章目录
- 一、问题定位
- 二、原理分析
- 三、解决方案与实践案例
- 步骤 1: 在 `model.train()` 调用中设置 `save_json=True`
- 步骤 2: 修改 `validator.py` 的 `__call__` 方法
- 四、总结与展望
- 五、参考文献
重要说明: 本篇博客内容基于 YOLOv12 进行演示,但其核心原理和配置方法适用于所有基于 Ultralytics 框架开发的 YOLO 版本。在应用代码修改时,请注意核对你所使用的具体 Ultralytics 库版本中的文件路径和代码行号,并确保已完成本系列教程前置步骤(特别是验证时启用 COCO API 的相关配置)。
一、问题定位
在使用 Ultralytics YOLO 进行模型训练时,标准的训练日志通常只展示损失函数(如 box_loss, cls_loss, dfl_loss)的变化。虽然这些损失值能反映训练趋势,但开发者往往更关心模型在验证集上的实际性能指标,特别是标准的 COCO 指标(如 mAP),以便:
- 实时评估模型效果: 了解每个 Epoch 后模型性能的真实提升。
- 指导超参数调优: 根据 mAP 的变化趋势调整学习率等参数。
- 实现基于 mAP 的早停 (Early Stopping): 在 mAP 不再提升时及时停止训练,节省资源。
- 选择最佳模型: 基于验证集上的最高 mAP 保存
best.pt
权重。
然而,默认的 model.train()
流程,即使配置了验证数据集,其在每个 Epoch 结束后的验证阶段,通常不直接计算和输出完整的 COCO 指标。本文旨在解决此问题,实现在训练过程中实时获取并监控 COCO mAP。
二、原理分析
Ultralytics YOLO 的训练流程 (model.train()
) 内部会周期性地调用验证器 (Validator
) 对模型在验证集上的性能进行评估。这个验证阶段与单独运行 model.val()
使用的是相似的逻辑,都涉及 ultralytics/engine/validator.py
(或其子类,如 ultralytics/models/yolo/detect/val.py
)。
要在训练的验证阶段计算 COCO 指标,需要克服以下关键点:
- 数据收集: 验证器在推理过程中需要收集所有预测框的信息,并将其整理成 COCO 格式的预测结果列表(通常存储在
self.jdict
变量中)。这需要save_json=True
参数被传递给验证器。 - 触发 COCO API 调用: 在验证器完成对验证集的推理后,需要显式调用能够处理
self.jdict
并与 Ground Truth JSON 文件交互,最终调用pycocotools
计算 COCO 指标的函数(如self.eval_json()
)。 - 结果整合: 计算出的 COCO 指标需要被整合回训练过程的日志和结果记录中,例如更新
results
字典,并可能影响fitness
值(用于判断是否为最佳模型)。
默认情况下,validator.py
在训练模式 (self.training = True
) 下执行 __call__
方法时,可能并未完整执行保存预测 JSON 文件并调用 self.eval_json()
的逻辑。因此,需要对 validator.py
的 __call__
方法进行修改,以确保在训练的验证阶段强制执行这一评估流程。
重要前提: 本文的修改依赖于 本系列教程前置步骤已正确完成。即,ultralytics/models/yolo/detect/val.py
(或对应检测模型的验证器)中关于 is_coco
的判断、anno_json
路径的定位、以及关键的 self.class_map
针对自定义数据集的调整必须已经生效。否则,即使触发了 eval_json()
,也可能因为找不到标注文件或类别映射错误而失败。
三、解决方案与实践案例
以下步骤展示了如何修改配置和代码,以在 model.train()
过程中启用实时的 COCO API 评估。
步骤 1: 在 model.train()
调用中设置 save_json=True
在启动训练时,需要将 save_json
参数设置为 True
。这将确保在每个 Epoch 结束后的验证阶段,验证器会收集预测结果并准备好生成 COCO 格式的 JSON 数据。
# --- 代码示例:调用 model.train() 并启用 save_json ---
from ultralytics import YOLO # 导入 YOLO 类
# 加载预训练模型或从 YAML 构建新模型
model = YOLO('yolov12n.pt') # 或 yolov12n.yaml
# 调用训练方法
results = model.train(
data='path/to/your/xxx.yaml', # 指向你的 COCO 格式数据集配置文件
# ... 其他训练参数 ...
save_json=True, # 关键:确保验证阶段收集预测结果以供 COCO API 使用
# ... 其他训练参数 ...
)
- 说明:
save_json=True
本身并不直接计算 COCO 指标,但它是后续步骤中触发 COCO API 调用的必要条件。
步骤 2: 修改 validator.py
的 __call__
方法
需要修改基础验证器类 ultralytics/engine/validator.py
中的 __call__
方法。在其 if self.training:
代码块内部,返回结果之前,添加强制保存预测 JSON 文件并调用 self.eval_json()
的逻辑。
- 代码位置 (仅供参考,请核对你的版本):
ultralytics/engine/validator.py
L205 附近,在if self.training:
块内,return {k: round(float(v), 5) for k, v in results.items()}
语句之前。
# --- validator.py __call__ 方法修改示例 ---
# ... 省略 __call__ 方法前面的代码 ...
if self.training: # 判断当前是否在训练模式下的验证阶段
model.float() # 确保模型在 float 模式
# ------------- 新添加的核心逻辑 ------------ #
if self.args.save_json and self.jdict: # 检查 save_json 是否为 True 且 self.jdict (预测结果列表) 是否有内容
pred_json_path = self.save_dir / "predictions.json"
LOGGER.info(f"Saving training validation predictions to {pred_json_path}...")
with open(str(pred_json_path), "w") as f:
json.dump(self.jdict, f) # 保存预测结果
LOGGER.info(f"Calculating COCO metrics for epoch validation...")
stats = self.eval_json(stats) # ★★★ 触发 COCO API 评估 ★★★
stats['fitness'] = stats['metrics/mAP50-95(B)'] # 更新 fitness 值,用于判断最佳模型 ('best.pt') ,这里默认使用 mAP50-95 作为 fitness
# 你可以根据需要调整 fitness 的计算方式,例如结合 mAP50
# stats['fitness'] = 0.1 * stats['metrics/mAP50(B)'] + 0.9 * stats['metrics/mAP50-95(B)']
# ------------- 新添加的核心逻辑结束 ------------ #
results = {**stats, **trainer.label_loss_items(self.loss.cpu() / len(self.dataloader), prefix="val")}
return {k: round(float(v), 5) for k, v in results.items()}
else:
# ... 省略后续代码 ...
- 修改说明: 这个修改的核心是在训练的验证环节,强制执行
predictions.json
的保存和eval_json()
的调用。eval_json()
会利用pycocotools
计算 COCO 指标,并将metrics/mAP50-95(B)
,metrics/mAP50(B)
等指标添加到stats
字典中。随后,这些指标会被包含在返回给训练器的results
中,并打印在训练日志里,同时fitness
值(通常基于 mAP)也会被正确更新,用于模型选择。
四、总结与展望
通过在调用 model.train()
时设置 save_json=True
,并对 ultralytics/engine/validator.py
中的 __call__
方法进行适当修改,我们成功地在 Ultralytics YOLO 的训练过程中实现了每个 Epoch 验证阶段的实时 COCO API 评估。这使得开发者能够:
- 实时监控关键性能指标 mAP 的变化。
- 更有效地进行模型调优和选择。
- 可能基于 mAP 实现更可靠的早停策略。
关键要点回顾:
- 训练启动时,设置
model.train(..., save_json=True, ...)
。 - 修改
validator.py
的__call__
方法,在if self.training:
分支内,当条件满足时调用self.eval_json(stats)
。 - 确保
fitness
值基于计算出的 COCO 指标进行更新。
希望本教程能帮助你更深入地理解和控制 Ultralytics YOLO 的训练评估流程,获取更精确实时的模型性能反馈。
五、参考文献
- Ultralytics YOLO 官方文档:
model.train()
API 文档: https://docs.ultralytics.com/modes/train/- Validator Class (可能在 API 文档或源码注释中): 查阅 Ultralytics GitHub 仓库中的
ultralytics/engine/validator.py
- 本系列教程前置文章 (重要依赖):
- 教程(3): 验证(Val)时启用 COCO API 评估 (确保
val.py
修改已完成)
- 教程(3): 验证(Val)时启用 COCO API 评估 (确保
- COCO Dataset 官网: https://cocodataset.org/#detection-eval
pycocotools
GitHub Repository: https://github.com/cocodataset/cocoapi