【YOLO】YOLOv8训练自定义数据集（4种方式）

YOLOv8 出来一段时间了&＃xff0c;继承了分类、检测、分割&＃xff0c;本文主要实现自定义的数据集&＃xff0c;使用 YOLOV8 进行检测模型的训练和使用

YOLOv8 此次将所有的配置参数全部解耦到配置文件 default.yaml&＃xff0c;不再类似于 YOLOv5&＃xff0c;一部分在配置文件&＃xff0c;一部分在 train.py 中

1. 运行环境

windows11 和 Ubuntu20.04&＃xff08;建议使用 Linux 系统&＃xff09;

首先切换到自己建立的虚拟环境安装 pytorch

torch 1.12.0&＃43;cu116&＃xff08;根据自身设备而定&＃xff09;
torchvision 0.13.0&＃43;cu116&＃xff08;根据自身设备而定&＃xff09;


安装完成后&＃xff0c;使用 git 命令将源码克隆下来

git clone https://github.com/ultralytics/ultralytics.git


参照官网&＃xff0c;直接使用以下语句即可导入项目所需要的库

pip install ultralytics


根据官方的解释&＃xff0c;pip 的 ultralytics 库包含了 requirements.txt中的所有库

2. 自定义数据集

我自己准备了一批 熊猫、老虎的图片作为实验数据集&＃xff0c;文件夹命名为 data &＃xff08;文件路径&＃xff1a;/home/mango/ultralytics/data&＃xff09;&＃xff0c;对数据集进行划分后的文件夹整体目录结构如下

images 下包含 train、val 文件夹&＃xff0c;这两个文件夹下包含此次需要的 图片信息
labels 下包含 train、val 文件夹&＃xff0c;这两个文件夹下包含此次需要的 对应图片的标注信息
test_images 下包含的是提供测试的数据集

ps&＃xff1a;数据集形式共有3种&＃xff0c;可根据自己实际情况采用

1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]


3. 模型训练&＃xff08;四种方式&＃xff09;

首先在 data 文件夹下新建一个数据加载配置文件 animal.yaml

train: /home/mango/ultralytics/data//images/train
val: /home/mango/ultralytics/data/images/val
# number of classes
nc: 2
# class names
names: [&＃39;panda&＃39;, &＃39;tiger&＃39;]


接下来就可以准备开始训练了
打开终端&＃xff0c;进入虚拟环境&＃xff0c;进入yolov8的文件夹&＃xff0c;
考虑到命令行模式下下载模型可能有点慢&＃xff0c;所以先在官方仓库下载好模型&＃xff0c;并放入新建的 weights &＃xff08;文件路径&＃xff1a;/home/mango/ultralytics/weights&＃xff09;目录下

YOLOv8模型仓库

(1) 第一种方式&＃xff08;参数重写&＃xff09;

参数很多&＃xff0c;建议查看 官网文档

下面是yolov8官方给定的命令行训练/预测/验证/导出方式&＃xff1a;

yolo task&＃61;detect mode&＃61;train model&＃61;yolov8n.pt args...classify predict yolov8n-cls.yaml args...segment val yolov8n-seg.yaml args...export yolov8n.pt format&＃61;onnx args...


最后输入以下命令即可开始训练&＃xff08;参数很多可以修改&＃xff0c;建议查看官网文档&＃xff0c;或者查看/home/mango/ultralytics/ultralytics/yolo/cfg下的 default.yaml 文件&＃xff09;

yolo task&＃61;detect mode&＃61;train model&＃61;weights/yolov8n.pt data&＃61;data/animal.yaml batch&＃61;16 epochs&＃61;50 imgsz&＃61;640 workers&＃61;16 device&＃61;0


可以看到已经成功开始训练&＃xff0c;运行生成的权重、混淆矩阵等信息存在于/home/mango/ultralytics/runs 下面

(2) 第二种方式&＃xff08;重写配置文件&＃xff09;

可以新建一个配置文件&＃xff0c;例如&＃xff1a;demo.yaml&＃xff0c;参数配置内容从 /home/mango/ultralytics/ultralytics/yolo/cfg/default.yaml 复制即可

或者使用命令行

yolo copy-cfg


它会自动生成一个 default_copy.yaml (目录地址&＃xff1a;/home/mango/ultralytics/default_copy.yaml)

截取的部分参数信息如下&＃xff1a;

# Ultralytics YOLO &＃x1f680;, GPL-3.0 license
# Default training settings and hyperparameters for medium-augmentation COCO trainingtask: detect # inference task, i.e. detect, segment, classify
mode: train # YOLO mode, i.e. train, val, predict, export# Train settings -------------------------------------------------------------------------------------------------------
model: # path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: # path to data file, i.e. coco128.yaml
epochs: 100 # number of epochs to train for
patience: 50 # epochs to wait for no observable improvement for early stopping of training
batch: 16 # number of images per batch (-1 for AutoBatch)
imgsz: 640 # size of input images as integer or w,h
save: True # save train checkpoints and predict results
save_period: -1 # Save checkpoint every x epochs (disabled if <1)
cache: False # True/ram, disk or False. Use cache for data loading
device: # device to run on, i.e. cuda device&＃61;0 or device&＃61;0,1,2,3 or device&＃61;cpu
workers: 8 # number of worker threads for data loading (per RANK if DDP)
project: # project name
name: # experiment name
...
...


根据自身需求&＃xff0c;修改相应参数&＃xff0c;例如&＃xff1a;修改 model、data、epochs、batch

model: weights/yolov8n.pt # path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: data/animal.yaml # path to data file, i.e. coco128.yaml
epochs: 20 # number of epochs to train for
batch: 8 # number of images per batch (-1 for AutoBatch)


然后在终端输入下列代码行命令即可开始训练

yolo cfg&＃61;default_copy.yaml


ps: 还可以使用 yolo cfg&＃61;default_copy.yaml imgsz&＃61;320 batch&＃61;8 的方式修改 imgz、batch 等参数信息

可以看到也已经成功开始训练&＃xff0c;运行生成的权重、混淆矩阵等信息存在于/home/mango/ultralytics/runs 下面

(3) 第三种方式&＃xff08;python命令&＃xff09;

在 Python 环境中直接使用

from ultralytics import YOLO# 加载模型
# model &＃61; YOLO("yolov8n.yaml") # 从头开始构建新模型
model &＃61; YOLO("weights/yolov8n.pt") # 加载预训练模型&＃xff08;推荐用于训练&＃xff09;# Use the model
results &＃61; model.train(data&＃61;"data/animal.yaml", epochs&＃61;20, batch&＃61;8) # 训练模型


或者创建一个 demo.py, 将上述代码拷贝到 demo.py&＃xff0c;然后调用 python demo.py 即可

可以看到也已经成功开始训练&＃xff0c;运行生成的权重、混淆矩阵等信息存在于/home/mango/ultralytics/runs 下面

(4) 第四种方式&＃xff08;python文件调用&＃xff09;

进入 /home/mango/ultralytics/ultralytics/yolo 目录下&＃xff0c;复制 animal.yaml、yolov8n.pt 到 /home/mango/ultralytics/ultralytics/yolo/v8/detect 目录&＃xff0c;
修改 /home/mango/ultralytics/ultralytics/yolo/cfg/default.yaml 的 model、data 路径及其他参数信息

# Ultralytics YOLO &＃x1f680;, GPL-3.0 license
# Default training settings and hyperparameters for medium-augmentation COCO trainingtask: detect # inference task, i.e. detect, segment, classify
mode: train # YOLO mode, i.e. train, val, predict, export# Train settings -------------------------------------------------------------------------------------------------------
model: yolov8n.pt # path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: animal.yaml # path to data file, i.e. coco128.yaml
epochs: 100 # number of epochs to train for
patience: 50 # epochs to wait for no observable improvement for early stopping of training
batch: 16 # number of images per batch (-1 for AutoBatch)
imgsz: 640 # size of input images as integer or w,h
save: True # save train checkpoints and predict results
save_period: -1 # Save checkpoint every x epochs (disabled if <1)
cache: False # True/ram, disk or False. Use cache for data loading
device: # device to run on, i.e. cuda device&＃61;0 or device&＃61;0,1,2,3 or device&＃61;cpu
workers: 8 # number of worker threads for data loading (per RANK if DDP)


运行 train.py

可以看到也已经成功开始训练&＃xff0c;运行生成的权重、混淆矩阵等信息存在于/home/mango/ultralytics/yolo/v8/detect/runs 下面

4. 模型预测

可修改的参数很多&＃xff0c;建议查看 官网文档

和模型训练一样&＃xff0c;预测同样可以采用不同的方式去实现&＃xff0c;这里展示其中一种方法&＃xff0c;主要目前还是看看模型效果
将训练得到的 best.pt 复制到 /home/mango/ultralytics/weights 下&＃xff0c;执行如下指令

yolo detect predict model&＃61;weights/best.pt source&＃61;data/test_images save&＃61;True


结果如下&＃xff1a;

最后一张图没有检测好&＃xff0c;估计是跟我训练数据集数量有关系&＃xff08;总共40&＃43;张&＃xff09;&＃xff0c;还有一个就是训练轮次&＃xff08;50轮&＃xff0c;花了不到2分钟时间&＃x1f602;&＃xff09;

&＃x1f44d; 但总体来说&＃xff0c;效果还是可以的&＃xff0c;速度精度都不低 &＃x1f525;

5. 最后

&＃x1f680; 接下来准备试试 onnx、和 TensorRT 的部署 ❗️