[YOLO V8]基于自定义数据集的简单训练和推理

2023 年 11 月 28 日

662 次浏览

3608字数

零、YOLO V8介绍

YOLO（You Only Look Once）系列模型在计算机视觉界名声大噪。YOLO 的名气归因于其在保持较小模型尺寸的同时具有相当高的准确性。YOLO 模型可以在单个 GPU 上进行训练，这使得开发人员可以廉价地使用它。机器学习从业者可以以低成本将其部署在边缘硬件或云端。

YOLOv8 是最新最先进的 YOLO 模型，可用于对象检测、图像分类和实例分割任务。YOLOv8 由 Ultralytics 开发，该公司还创建了具有影响力和行业定义的 YOLOv5 模型。与 YOLOv5 相比，YOLOv8 包含大量架构和开发人员体验更改和改进。

一、安装YOLO V8

可以选择直接在虚拟环境里使用pip进行安装 pip install ultralytics

或者clone官方仓库进行安装 git clone https://github.com/ultralytics/ultralytics

二、数据准备

这里使用的是水下机器人结构与控制课上采集的，三个不同颜色球体的数据，我们已经做好了数据的标注工作，数据标注我会另外写一下，所以在这里先省略（绝对不是懒了）

YOLO训练前我们需要的的文件结构：

├─data
│  │  ball.yaml
│  │  
│  ├─test
│  │  ├─images
│  │  │      298.jpg
│  │  │      
│  │  └─labels
│  │          298.txt
│  │          
│  ├─train
│  │  │  labels.cache
│  │  │  
│  │  ├─images
│  │  │      283.jpg
│  │  │      284.jpg
│  │  │      285.jpg
│  │  │      286.jpg
│  │  │      287.jpg
│  │  │      
│  │  └─labels
│  │          283.txt
│  │          284.txt
│  │          285.txt
│  │          286.txt
│  │          287.txt
│  │          
│  └─val
│      │  labels.cache
│      │  
│      ├─images
│      │      13.jpg
│      │      14.jpg
│      │      15.jpg
│      │      
│      └─labels
│              13.txt
│              14.txt
│              15.txt

ball.yaml记录了我们训练时所需要的数据源和分类类别，例如：

train: **/data/train # 训练集文件夹
val: **/data/val # 验证集文件夹
test: **/data/test # 测试集文件夹

# 分类的数量
nc: 3

# 每个类的名字
names: ['red_ball', 'yellow_ball', 'green_ball']

剩余三个文件夹，train、val、test下各有两个文件夹，分别为images和labels，分别存放图片数据和标记好的标签数据。

三、训练模型

这里我使用的是Python导入库的方式进行运行，脚本如下：

from ultralytics import YOLO

if __name__ == '__main__': # 请将代码放在__main__下运行，否则会报错
    model = YOLO('yolov8n.pt') # 程序会自动从Github下载默认的权重文件，你也可以提前下载好后放在这里
    model.train(data='**/data/ball.yaml', epochs=100, batch=8, amp=False) # 开始训练

由于我的显卡显存只有4G，所以这里的batch调的比较小，如果显存够大，可以再往上调整一些，训练效果会更好。

同时，YOLO V8也支持基于命令行进行训练，你可以用如下的命令替代这个脚本：

yolo task=detect mode=train model=yolov8n.pt data='**/data/ball.yaml' epochs=100 batch=8 amp=False

四、预测

直接使用命令行：

yolo task=detect mode=predict model=<你的模型位置> source=<待预测的数据> device=0

或者使用Python脚本

from ultralytics import YOLO

if __name__ == '__main__': # 请将代码放在__main__下运行，否则会报错
    # 加载训练好的模型
    model = YOLO('<训练好的权重文件>')

    # 利用模型进行预测
    results = model('<你要预测的文件>', save=True)

附录、可能遇到的一些问题

1.有可能会默认使用CPU进行训练。

可以使用下列的代码检查CUDA是否安装正确，以及确认所安装的Torch版本是否和CUDA是配套的。

import torch
#查看版本
print(torch.__version__)
#查看gpu是否可用
print(torch.cuda.is_available())
#返回设备gpu个数
print(torch.cuda.device_count())
# 查看对应CUDA的版本号
print(torch.backends.cudnn.version())
print(torch.version.cuda)
#退出python
quit()

另外，利用pip install ultralytics执行安装时，有可能默认安装的是CPU版本的Torch，需要到Torch官网找对应的版本覆盖安装。

Torch官网：https://pytorch.org/

过往版本Torch:https://pytorch.org/get-started/previous-versions/