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

原创

用户9029617

修改于 2024-06-01 18:05:23

1290

修改于 2024-06-01 18:05:23

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_callback


if __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 YOLO

if __name__ == "__main__":
    model = YOLO("yolov8n.pt")

    model.train(
        data="./coco128.yaml",
        epochs=3, 
        imgsz=320,
        # 开启DDP
        device=[0,1],
    )

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

def add_integration_callbacks(instance):
    ...
    
    # Load training callbacks
    if "Trainer" in instance.__class__.__name__:
        from .clearml import callbacks as clear_cb
        from .comet import callbacks as comet_cb
        from .dvc import callbacks as dvc_cb
        from .mlflow import callbacks as mlflow_cb
        from .neptune import callbacks as neptune_cb
        from .raytune import callbacks as tune_cb
        from .tensorboard import callbacks as tb_cb
        from .wb import callbacks as wb_cb

        from swanlab.integration.ultralytics import return_swanlab_callback
        sw_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了。undefined项目名也可以由model.train()的project参数定义，实验名由name参数定义。

2.2 代码案例

from ultralytics import YOLO

if __name__ == "__main__":
    model = YOLO("yolov8n.pt")

    model.train(
        data="./coco128.yaml",
        epochs=3, 
        imgsz=320,
        # 开启DDP
        device=[0,1,2,3],
        # 可以通过project参数设置SwanLab的project，name参数设置SwanLab的experiment_name
        project="YOLO_COCO128",
        name="yolov8n"
    )