Ultralytics x SwanLab：可视化YOLO模型训练

Ultralytics是YOLO官方团队推出的CV训练与推理框架，不仅支持目标检测任务，还支持分割、姿态识别、分类等更多任务。

在这里插入图片描述

SwanLab是一个深度学习实验管理与训练可视化工具，由西安电子科技大学团队打造，融合了Weights & Biases与Tensorboard的特点，能够方便地进行训练可视化、多实验对比、超参数记录、大型实验管理和团队协作，并支持用网页链接的方式分享你的实验。

在这里插入图片描述

你可以使用Ultralytics快速进行目标检测模型训练，同时使用SwanLab进行实验跟踪与可视化。可视化结果(可以在这里直接预览。)

集成SwanLab有两种方式：

使用add_swanlab_callback函数：无需修改源码，仅适用于单卡训练场景
使用return_swanlab_callback函数：需要修改源码，适用于单卡以及多卡DDP训练场景

下面将用一个yolo在coco128数据集上的训练，来介绍如何将Ultralytics与SwanLab配合起来。

1.1 引入add_swanlab_callback

from swanlab.integration.ultralytics import add_swanlab_callback

SwanLabCallback是适配于Transformers的日志记录类。

SwanLabCallback可以定义的参数有：

project、experiment_name、description 等与 swanlab.init 效果一致的参数, 用于SwanLab项目的初始化。
你也可以在外部通过swanlab.init创建项目，集成会将实验记录到你在外部创建的项目中。

1.2 引入add_swanlab_callback

下面是使用yolov8n模型在coco数据集上的训练，只需将model传入add_swanlab_callback函数，即可完成与SwanLab的集成。

from ultralytics import YOLO
from swanlab.integration.ultralytics import add_swanlab_callbackif __name__ == "__main__":model = YOLO("yolov8n.pt")# 添加swanlab回调add_swanlab_callback(model)model.train(data="./coco128.yaml",epochs=3, imgsz=320,)

如果需要自定义SwanLab的项目、实验名等参数，则可以在add_swanlab_callback中添加：

add_swanlab_callback(model,project="YOLO_COCO128",experiment_name="yolov8n",description="yolov8n在coco128数据集上的训练。",)

然后开始训练即可！

2.1 多卡训练/DDP训练

swanlab>=0.3.8

在Ultralytics多卡训练的场景下，由于启动训练的方式与单卡完全不同，所以需要用一种不同的方式接入SwanLab回调。

这是一个ultralytics开启DDP训练的样例代码：

from ultralytics import YOLOif __name__ == "__main__":model = YOLO("yolov8n.pt")model.train(data="./coco128.yaml",epochs=3, imgsz=320,# 开启DDPdevice=[0,1],)

我们需要修改ultralytics的源码，去到ultralytics/utils/callbacks/base.py，找到add_integration_callbacks函数，添加下面的三行代码：

def add_integration_callbacks(instance):...# Load training callbacksif "Trainer" in instance.__class__.__name__:from .clearml import callbacks as clear_cbfrom .comet import callbacks as comet_cbfrom .dvc import callbacks as dvc_cbfrom .mlflow import callbacks as mlflow_cbfrom .neptune import callbacks as neptune_cbfrom .raytune import callbacks as tune_cbfrom .tensorboard import callbacks as tb_cbfrom .wb import callbacks as wb_cbfrom swanlab.integration.ultralytics import return_swanlab_callbacksw_cb = return_swanlab_callback()callbacks_list.extend([..., sw_cb])

然后运行，就可以在ddp下正常跟踪实验了。

如果需要自定义SwanLab的项目、实验名等参数，则可以在return_swanlab_callback中添加：

return_swanlab_callback(model,project="ultralytics",experiment_name="yolov8n",description="yolov8n在coco128数据集上的训练。",)

写入源码之后，之后运行就不需要在训练脚本中增加add_swanlab_callback了。
项目名也可以由model.train()的project参数定义，实验名由name参数定义。

2.2 代码案例

from ultralytics import YOLOif __name__ == "__main__":model = YOLO("yolov8n.pt")model.train(data="./coco128.yaml",epochs=3, imgsz=320,# 开启DDPdevice=[0,1,2,3],# 可以通过project参数设置SwanLab的project，name参数设置SwanLab的experiment_nameproject="YOLO_COCO128",name="yolov8n")