使用 yolov11n 为基底模型训练自己的模型

442 阅读1分钟

数据集

在项目目录下创建 datasets 目录, 再在此目录下创建数据集名称如 mcc; 在此目录下按下列层级创建文件:

  • images
    • train (训练集图片)
      • 001.jpg
      • 002.jpg
      • 003.jpg
      • 004.jpg
      • 005.jpg
    • val (验证集图片)
      • 101.jpg
      • 102.jpg
  • labels
    • classes.txt
    • train (训练集标签文件,与训练集图片一一对应)
      • 001.txt
      • 002.txt
      • 003.txt
      • 004.txt
      • 005.txt
    • val
      • 101.txt
      • 102.txt

训练表述文件

命名为 mcc.yaml

path: ./mcc # dataset root dir
train: images/train/ # train images (relative to 'path') 39384 images
val: images/val/ # val images (relative to 'path') 15062 images
test: # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview

# Classes
names:
  0: person
  1: bicycle
  2: car

模型训练

  1. 命令行方式

yolo detect train model=./yolo11n.pt data="mcc.yaml" epochs=50 imgsz=640

此处使用牵引学习,需要先有 yolov11n.pt 模型作为基础模型

若本地没有yolov11n.pt模型,则可以使用下述代码,运行后会自动下载yolov11n.pt模型

from ultralytics import YOLO

if __name__ == "__main__":
    model = YOLO("yolo11n.pt") # https://github.com/ultralytics/assets/releases/tag/v8.3.0
    img = "catdog.jpg"
    results = model(img)
    for result in results:
        print(result.boxes)  # Print detection boxes
        result.show()  # Display the annotated image
  1. 代码方式
from ultralytics import YOLO

# 加载模型
model = YOLO("yolo11n.pt")

# 训练模型
train_results = model.train(
    data="mcc.yaml",  # 数据集 YAML 路径
    epochs=50,  # 训练轮次
    imgsz=640,  # 训练图像尺寸
    device="cpu",  # 运行设备,例如 device=0 或 device=0,1,2,3 或 device=cpu
)
  1. cfg方式

yolo copy-cfg

# default-copy.yaml
task: detect # (str) YOLO task, i.e. detect, segment, classify, pose, obb
mode: train # (str) YOLO mode, i.e. train, val, predict, export, track, benchmark

# Train settings ---------------------------
model: yolov11n.pt # (str, optional) path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: mcc.yaml # (str, optional) path to data file, i.e. coco8.yaml
epochs: 100 # (int) number of epochs to train for
...

yolo cfg=default-copy.yaml

验证模型

yolo detect predict model=runs/detect/train/weigths/best.pt source=./mcc.mp4 show=True


再次训练时需删除 ultralytics 目录下的 settings.yaml