YOLOv8入门-训练TT100K数据集实践

前言

本篇文章主要讲解yolov8模型训练自己数据集全过程，包括环境搭建、数据集准备、模型训练、模型验证、模型预测等

yolov8源码地址：https://github.com/ultralytics/ultralytics

一、环境搭建

在这之前，需要准备基础环境：
ubuntu20.04、python=3.8、pytorch=2.0.0、cuda=11.8

方法一：准备好基础环境后，进行自己创建的虚拟环境，安装ultralytics库，只需安装这一个

pip install ultralytics

方法二：先克隆git仓库，再进行安装

git clone https://github.com/ultralytics/ultralytics
cd ultralytics
pip install -e .

二、训练准备

第一步：新建datasets目录，将数据集放入，数据集结构如下：

第二步：在datasets目录下，新建一个data_tt100k.yaml文件

以下是我新建的，里面写绝对路径（主要是怕出错）：

数据集的准备已经就绪，索引文件在data目录下的train.txt/val.txt/test.txt中

第三步：进入ultralytics/ultralytics/cfg/models/v8文件夹下，修改网络模型配置文件，更改类别

三、开始训练

yolov8关于模型的各种参数其实都写到了一起（包括训练、验证和推理的参数），在/root/ultralytics/ultralytics/cfg/default.yaml中，通过使用这些指令我们可以实现各种所需的操作。

由于yolov8代码中没有提供模型训练代码，这里新建一个train.py文件，主要包括训练参数，代码如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO(r'yolov8.yaml')  # 不使用预训练权重训练
    model = YOLO(r'yolov8.yaml').load("yolov8n.pt")  # 使用预训练权重训练
    # 训练参数 ----------------------------------------------------------------------------------------------
    model.train(
        data=r'datasets/data_tt100k.yaml',
        epochs=300,  # (int) 训练的周期数
        patience=50,  # (int) 等待无明显改善以进行早期停止的周期数
        batch=32,  # (int) 每批次的图像数量（-1 为自动批处理）
        imgsz=640,  # (int) 输入图像的大小，整数或w，h
        save=True,  # (bool) 保存训练检查点和预测结果
        save_period=-1,  # (int) 每x周期保存检查点（如果小于1则禁用）
        cache=False,  # (bool) True/ram、磁盘或False。使用缓存加载数据
        device='',  # (int | str | list, optional) 运行的设备，例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
        workers=8,  # (int) 数据加载的工作线程数（每个DDP进程）
        project='runs/train',  # (str, optional) 项目名称
        name='exp',  # (str, optional) 实验名称，结果保存在'project/name'目录下
        exist_ok=False,  # (bool) 是否覆盖现有实验
        pretrained=True,  # (bool | str) 是否使用预训练模型（bool），或从中加载权重的模型（str）
        optimizer='SGD',  # (str) 要使用的优化器，选择=[SGD，Adam，Adamax，AdamW，NAdam，RAdam，RMSProp，auto]
        verbose=True,  # (bool) 是否打印详细输出
        seed=0,  # (int) 用于可重复性的随机种子
        deterministic=True,  # (bool) 是否启用确定性模式
        single_cls=False,  # (bool) 将多类数据训练为单类
        rect=False,  # (bool) 如果mode='train'，则进行矩形训练，如果mode='val'，则进行矩形验证
        cos_lr=False,  # (bool) 使用余弦学习率调度器
        close_mosaic=0,  # (int) 在最后几个周期禁用马赛克增强
        resume=False,  # (bool) 从上一个检查点恢复训练
        amp=True,  # (bool) 自动混合精度（AMP）训练，选择=[True, False]，True运行AMP检查
        fraction=1.0,  # (float) 要训练的数据集分数（默认为1.0，训练集中的所有图像）
        profile=False,  # (bool) 在训练期间为记录器启用ONNX和TensorRT速度
        freeze=None,  # (int | list, 可选) 在训练期间冻结前 n 层，或冻结层索引列表。
        # 分割
        overlap_mask=True,  # (bool) 训练期间是否应重叠掩码（仅适用于分割训练）
        mask_ratio=4,  # (int) 掩码降采样比例（仅适用于分割训练）
        # 分类
        dropout=0.0,  # (float) 使用丢弃正则化（仅适用于分类训练）
        # 超参数 ----------------------------------------------------------------------------------------------
        lr0=0.01,  # (float) 初始学习率（例如，SGD=1E-2，Adam=1E-3）
        lrf=0.01,  # (float) 最终学习率（lr0 * lrf）
        momentum=0.937,  # (float) SGD动量/Adam beta1
        weight_decay=0.0005,  # (float) 优化器权重衰减 5e-4
        warmup_epochs=3.0,  # (float) 预热周期（分数可用）
        warmup_momentum=0.8,  # (float) 预热初始动量
        warmup_bias_lr=0.1,  # (float) 预热初始偏置学习率
        box=7.5,  # (float) 盒损失增益
        cls=0.5,  # (float) 类别损失增益（与像素比例）
        dfl=1.5,  # (float) dfl损失增益
        pose=12.0,  # (float) 姿势损失增益
        kobj=1.0,  # (float) 关键点对象损失增益
        label_smoothing=0.0,  # (float) 标签平滑（分数）
        nbs=64,  # (int) 名义批量大小
        hsv_h=0.015,  # (float) 图像HSV-Hue增强（分数）
        hsv_s=0.7,  # (float) 图像HSV-Saturation增强（分数）
        hsv_v=0.4,  # (float) 图像HSV-Value增强（分数）
        degrees=0.0,  # (float) 图像旋转（+/- deg）
        translate=0.1,  # (float) 图像平移（+/- 分数）
        scale=0.5,  # (float) 图像缩放（+/- 增益）
        shear=0.0,  # (float) 图像剪切（+/- deg）
        perspective=0.0,  # (float) 图像透视（+/- 分数），范围为0-0.001
        flipud=0.0,  # (float) 图像上下翻转（概率）
        fliplr=0.5,  # (float) 图像左右翻转（概率）
        mosaic=1.0,  # (float) 图像马赛克（概率）
        mixup=0.0,  # (float) 图像混合（概率）
        copy_paste=0.0,  # (float) 分割复制-粘贴（概率）
    )

训练时比较重要的几个参数:

YOLO(): 这个参数里面可以写你模型 yaml文件的路径，也可以直接写 .pt 文件
YOLo().load(): 如果你不加 .load()，那就是不使用权重从头训练，加了就是使用使用预训练权重训练
data :数据集 yaml 文件的路径
epochs: 训练轮数
batch :batch size的大小
cache ：设置成 True 可以加速训练
patience :早停轮数，不想设置早停的话，可以把这个值设置成一个非常大的数
optimizer: 优化器，最新版应该是提供了 7 种，注释有写包含哪几种，还可以设置成 auto
close_mosaic: 最后多少轮关闭马赛克数据增强，这是一种训练技巧，很有效，一般设置 10-30
resume : 断点续训，一个最实用的参数，直接在这里写你断了的那轮的last.pt 的路径就接上了
single_cls: 数据集是单个类别的时候记得开启，不开其实也没事
cos_lr:是否开启余弦学习率
label_smoothing : 数据集标注质量很差的时候可以考虑给这个参数设置个0.1-0.5
device :多 GPu 训练时这里直接写一个列表，比如[0，1]

训练结果保存在runs文件夹下：

PR曲线图：

四、模型验证

YOLO(): 这里写你想要验证的权重的路径，如best.pt
data :这里写自己数据集 yaml 文件的路径
split :写val就是看验证集的指标，写 test 就是看测试集的指(前提是划分了测试集)
batch ：测试速度时一般都设置成1（验证比较稳定），设置越大，速度越快

同上，新建val.py文件，代码如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO('/root/ultralytics/runs/detect/train2/weights/best.pt')  
    # 验证模型
    metrics=model.val(
        val=True,  # (bool) 在训练期间进行验证/测试
        data='/root/ultralytics/datasets/data_tt100k.yaml',
        split='val',  # (str) 用于验证的数据集拆分，例如'val'、'test'或'train'
        batch=1,  # (int) 每批的图像数量（-1 为自动批处理）
        imgsz=640,  # 输入图像的大小，可以是整数或w，h
        device='0',  # 运行的设备，例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
        workers=8,  # 数据加载的工作线程数（每个DDP进程）
        save_json=False,  # 保存结果到JSON文件
        save_hybrid=False,  # 保存标签的混合版本（标签 + 额外的预测）
        conf=0.001,  # 检测的目标置信度阈值（默认为0.25用于预测，0.001用于验证）
        iou=0.6,  # 非极大值抑制 (NMS) 的交并比 (IoU) 阈值
        project='runs/val',  # 项目名称（可选）
        name='exp',  # 实验名称，结果保存在'project/name'目录下（可选）
        max_det=300,  # 每张图像的最大检测数
        half=False,  # 使用半精度 (FP16)
        dnn=False,  # 使用OpenCV DNN进行ONNX推断
        plots=True,  # 在训练/验证期间保存图像
    )

    print(f"mAP50-95: {metrics.box.map}") # map50-95
    print(f"mAP50: {metrics.box.map50}")  # map50
    print(f"mAP75: {metrics.box.map75}")  # map75
    speed_metrics = metrics.speed
    total_time = sum(speed_metrics.values())
    fps = 1000 / total_time
    print(f"FPS: {fps}") # FPS

验证结果如下：

五、模型推理

新建detect.py文件，代码如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO(r'runs/detect/train/weights/best.pt')  # YOLOv8n模型
    model.predict(
        source=r'datasets/images/test',
        save=True,  # 保存预测结果
        imgsz=640,  # 输入图像的大小，可以是整数或w，h
        conf=0.25,  # 用于检测的目标置信度阈值（默认为0.25，用于预测，0.001用于验证）
        iou=0.45,  # 非极大值抑制 (NMS) 的交并比 (IoU) 阈值
        show=False,  # 如果可能的话，显示结果
        project='runs/predict',  # 项目名称（可选）
        name='exp',  # 实验名称，结果保存在'project/name'目录下（可选）
        save_txt=False,  # 保存结果为 .txt 文件
        save_conf=True,  # 保存结果和置信度分数
        save_crop=False,  # 保存裁剪后的图像和结果
        show_labels=True,  # 在图中显示目标标签
        show_conf=True,  # 在图中显示目标置信度分数
        vid_stride=1,  # 视频帧率步长
        line_width=3,  # 边界框线条粗细（像素）
        visualize=False,  # 可视化模型特征
        augment=False,  # 对预测源应用图像增强
        agnostic_nms=False,  # 类别无关的NMS
        retina_masks=False,  # 使用高分辨率的分割掩码
        boxes=True,  # 在分割预测中显示边界框
    )

六、onnx导出

方法一：CLI 指令导出方式

yolo task=detect mode=export model=runs/detect/train/weights/best.pt

方法二：python 指令导出方式

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n.pt')  # load an official model
model = YOLO('path/to/best.pt')  # load a custom trained

# Export the model
model.export(format='onnx')

导出有关的具体参数如下：

format torchscript 导出的模型格式
keras False 是否使用Keras
optimize False TorchScript: 是否针对移动设备进行优化
int8 False CoreML/TF的INT8量化
dynamic False ONNX/TF/TensorRT: 是否动态轴
simplify False ONNX: 是否简化模型
opset ONNX: 操作版本（可选）
workspace 4 TensorRT: 工作空间大小（GB）
nms False CoreML: 是否添加NMS

注释：TorchScripts是pytorch的模型导出工具，INT8(8位整数量化)是一种量化方法，可将神经网络参数表示为8位整数，以降低存储的计算成本。onnx是一种跨平台、开放式的机器学习框架。TensorRT是一种用于加速深度学习推理的高性能引擎，CoreML是苹果公司推出的机器学习框架。Keras是一种流行的深度学习框架。

使用Netron对onnx网络结构查看