机器视觉-yolov8集成wandb进行训练监控

Wandb 概述

WandB 即 Weigtht and Bias的缩写, 是深度学习中经常使用的metrics记录工具, 功能比TensorBoard强大, 主要功能有:

• 它可以记录每次训练的版本信息, 包括超参、tag、project 等
• 自动上传云端，方便将多台机器的实验做对比分析
• 强大的表格功能，可以轻松管理海量模型
• 支持离线数据上传功能，实验完成后可以离线上传数据
• 支持局域网部署服务器

官网注册账号

官网: https://wandb.ai/site
我使用 google 账号注册了一个wandb的账号.
注册后, 可以从 https://wandb.ai/quickstart 页面看到自己的 Api key.

安装并登录

#  安装
pip install wandb

# 登录命令, 登录完成后会将 api_key 记录到 C:\Users\YOUR_USER\.netrc 文件中. 
wandb.exe login

yolov8设置wandb离线模式

如果访问网络wandb官网的网络不好, 会拉慢整个训练速度, 所以可以设置成离线模式, 待训练完成后在手工上报metrics数据.

from wandb.integration.ultralytics import add_wandb_callback
import wandb

os.environ["WANDB_API_KEY"]="WANDB_API_KEY"
os.environ["WANDB_MODE"]="offline" # online or offline

在完成 yolov8 一个experiment之后, console 会智能提示一个上传离线wandb日志的命令, 比如:

wandb sync D:\my_workspace\source\yolo8\wandb\offline-run-20240216_193023-bpqnk0yg

YoloV8集成 wangdb 的代码示例

个人认为predict无需进行监控, 监控training阶段会更有意义

import os
from IPython import display
import ultralytics
from ultralytics import YOLO, settings
from os import path
from wandb.integration.ultralytics import add_wandb_callback
import wandb

os.environ["WANDB_API_KEY"]="your_api_key"
os.environ["WANDB_MODE"]="online" # offline or online

def train():
    # 修复 C:\Users\MY_USER\AppData\Roaming\Ultralytics\settings.yaml 文件中的图像目录
    image_path = r'''C:\Users\dorothy\AppData\Roaming\Python\Python311\Scripts\datasets'''
    settings.update({"datasets_dir": image_path})

    project_name ="yolov8_test"
    run_name ="yolov8n_name114"
    epochs_num=1
    batch_num=8

    # init wandb
    hyperparameters = dict(
        epochs=epochs_num,
        batch_size=batch_num
    )
    config_dictionary = dict(
        yaml=r"D:\my_workspace\py_code\yolo8\Lib\site-packages\ultralytics\cfg\default.yaml",
        params=hyperparameters,
    )
    config_dictionary["project"]=project_name
    config_dictionary["name"] = run_name
    config_dictionary["notes"] ="my first test"
    config_dictionary["tags"] =["baseline", "paper"]
    config_dictionary["job_type"] = "training" # training or inference
    run=wandb.init(config=config_dictionary)

    # 可提前将预训练的yolov8n.pt文件下载到本脚本同一目录下,地址 https://github.com/ultralytics/assets/releases
    model = YOLO("yolov8n.pt")

    #Add W&B Callback for Ultralytics, 设置callback后, 在后续的train/val过程将会自动将指标数据传到Wandb中
    add_wandb_callback(model, enable_model_checkpointing=True)

    # 设置训练data yaml文件
    dataset_yaml = r"C:\Users\dorothy\AppData\Roaming\Python\Python311\site-packages\ultralytics\cfg\datasets\coco8.yaml"

    # 在本脚本同一目录下, 会自动生成 runs 结果目录
    results = model.train(project=project_name, name=run_name, data=dataset_yaml, epochs=epochs_num, imgsz=640, device="cpu", batch=batch_num, seed=1)

    #Finalize the W&B Run, 告知wandb dashboard将会会标识train/val过程已经完成
    wandb.finish()


def predict():
    model = YOLO(r"D:\my_workspace\source\yolo8\runs\detect\train9\weights\last.pt")
    add_wandb_callback(model, enable_model_checkpointing=True)
    image_file1 = r"D:\my_workspace\py_code\yolo8\Lib\site-packages\ultralytics\assets\bus.jpg"
    image_file2 = r"D:\my_workspace\py_code\yolo8\Lib\site-packages\ultralytics\assets\zidane.jpg"
    results_list = model.predict(source=[image_file1, image_file2], show=False, save=True, save_conf=True,
                                 save_txt=True)
    for results in results_list:
        boxes = results.boxes
        speed = results.speed
        names = results.names
        json = results.tojson()
        image_path = results.path
        print("====")
        print(image_path)
        print(names)
        print(json)


if __name__ == '__main__':
    display.clear_output()
    ultralytics.checks()
    train()

参考

https://docs.ultralytics.com/integrations/weights-biases/#key-features-of-the-weights-biases-dashboard
wandb使用教程(四)：wandb本地部署 https://zhuanlan.zhihu.com/p/521663928?utm_id=0https://zhuanlan.zhihu.com/p/342300434?utm_id=0https://blog.csdn.net/qq_42312574/article/details/132662921https://blog.csdn.net/crasher123/article/details/132644320