一文读懂yolo11模型训练

一、环境准备

Anaconda安装
简介

Anaconda 是一个流行的开源 Python 发行版,专注于数据科学、机器学习、科学计算和分析等领域。它提供了一个强大的包管理器环境管理器,名为 Conda,以及一个预装了大量科学计算和数据分析相关库的环境。

简单说就是可以创建N个互不干扰的环境,避免包产生冲突,用哪个就激活哪个。

下载安装

​ 地址:https://www.anaconda.com/download/success

下载安装包,一路下一步即可,注意最后要勾选添加环境变量不然后面还要手动添加,我这边装过了就不重装了。

主要命令
# 查看已创建的环境
conda env list
# 指定python版本创建环境
conda create -n envName python=version
# 激活指定环境
conda activate xxx
# 删除指定环境
conda env remove --name xxx
# 查看当前环境安装包
conda list

anaconda在安装时会默认创建一个base环境,没有通过activate激活环境时,默认使用的是base环境。

Pytorch安装

​ 在安装之前先确定电脑有没有英伟达的显卡,打开任务管理器-性能,翻到下面即可查看显卡型号。

请添加图片描述

如果没有相关的显卡的话安装Pytorch时只能安装cpu-only的包。

如果是英伟达的显卡,执行以下命令,查看CUDA版本

nvidia-smi

请添加图片描述

我们在安装选择版本时只要低于这个版本即可,并不是只能安装这个版本的包。如果说是英伟达的显卡但是这边输入命令没有返回时,可能是没装驱动,看下电脑有没有安装下面这个图标的软件,打开下载更新驱动。

请添加图片描述

创建环境

使用anaconda创建一个python3.10版本的隔离环境来使用

conda create -n yolo11 python=3.10

请添加图片描述

创建之后这边按照提示激活环境,激活后前面会显示环境名称。

conda activate yolo11

请添加图片描述

安装测试

​ Pytoch官网: https://pytorch.org/

打开Pytorch官网查看我们能安装的版本

请添加图片描述

如果CUDA低于11.8的话点击下面的历史版本,这边安装pytorch的时候尽量不要换源,用国内源的话很可能会下到CPU的版本。

因为我们后面要使用这个环境来训练yolo11的模型,所以还要考虑yolo11需要用到的版本,yolo11限制torch>=1.8.0且不等于2.4.0

我这边安装2.2.1的版本,记得要激活环境,在创建的环境里安装。

conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=12.1 -c pytorch -c nvidia

安装之后在输入python,然后输入import torch正常引入依赖即安装成功,可能会报一个NumPy版本高的错误,我们后续解决。如果是使用GPU训练,还要输入以下命令确认下是否能正常连接显卡驱动,如果返回True则正常,返回False可能是驱动版本问题,也可能是装了cpu的版本。

torch.cuda.is_available()
依赖安装

​ 项目地址: https://github.com/ultralytics/ultralytics

接下来安装训练模型所需依赖。将仓库中所需依赖放到requirements.txt文件中,torch手动安装了,所以将其注释掉。

certifi==2024.8.30
charset-normalizer==3.3.2
colorama==0.4.6
contourpy==1.3.0
cycler==0.12.1
dill==0.3.9
filelock==3.13.1
fonttools==4.54.1
fsspec==2024.2.0
idna==3.10
Jinja2==3.1.3
kiwisolver==1.4.7
MarkupSafe==2.1.5
matplotlib==3.9.2
mpmath==1.3.0
networkx==3.2.1
numpy==1.26.3
opencv-python==4.10.0.84
packaging==24.1
pandas==2.2.3
pillow==10.2.0
psutil==6.0.0
py-cpuinfo==9.0.0
pyparsing==3.1.4
python-dateutil==2.9.0.post0
pytz==2024.2
PyYAML==6.0.2
requests==2.32.3
scipy==1.14.1
seaborn==0.13.2
six==1.16.0
sympy==1.12
# torch==2.2.1+cu121
# torchaudio==2.2.1+cu121
# torchvision==0.17.1+cu121
tqdm==4.66.5
typing_extensions==4.9.0
tzdata==2024.2
ultralytics==8.3.0
ultralytics-thop==2.0.8
urllib3==2.2.3
onnx>=1.12.0
onnxslim==0.1.34
onnxruntime

安装所需依赖,这边可以用国内镜像源,下载会快些。

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

二、数据集准备

数据标注

数据标注可以分为自动标注和手动标注,自动的限制比较大,很多数据标注有问题,因为自动标注其实也是基于一个模型能力,如果你这个模型都能百分百标注我的对象了,那我还训练干啥,直接用这个模型不就好了,所以这边讲一下手动的标注方式。

Labelme使用

​ 项目地址: https://github.com/wkentaro/labelme

在仓库Release中下载软件并打开。首先在设置中选择自动保存和取消勾选同时保存图像数据,点击更改输出路径选择好标注数据的保存位置。如果每张图片标注的对象大小一致,则勾选编辑中的保留最后的标注。这样在标注下一张图片只需要拖动创建的矩形框就好。

请添加图片描述

点击打开目录,选择要标注的图像目录。右键创建矩形即可开始框选我们要标注的物体,然后在弹出的框中输入分类

请添加图片描述

右下角会显示我们标注的文件,打上勾的是已经标注过的,没有打勾的是未标注的。标注好后按下D键即可跳转下一张图片,按A返回上一张。

json转yolo

标注之后的数据是json格式,但是yolo使用的是txt格式并且存储的数据也有所区分,yolo格式如下:

<class> <x_center> <y_center> <width> <height>
数据分别表示 类别编号 边界框中心点的x坐标 边界框中心点的y坐标 边界框的宽度 边界框的高度 

所以要将我们标注的数据进行一些处理,网上很多,下面直接放源码。

import json
import os
import glob
from pathlib import Path

def convert_labelme_to_yolo(json_path, class_name_to_id):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    img_height = data['imageHeight']
    img_width = data['imageWidth']
    
    yolo_annotations = []

    for shape in data['shapes']:
        class_name = shape['label']
        points = shape['points']

        if class_name not in class_name_to_id:
            print(f"Warning: class {class_name} not found in mapping, skipping...")
            continue

        class_id = class_name_to_id[class_name]

        x_coords = [p[0] for p in points]
        y_coords = [p[1] for p in points]

        x_min = min(x_coords)
        y_min = min(y_coords)
        x_max = max(x_coords)
        y_max = max(y_coords)

        x_center = ((x_min + x_max) / 2) / img_width
        y_center = ((y_min + y_max) / 2) / img_height
        width = (x_max - x_min) / img_width
        height = (y_max - y_min) / img_height

        x_center = min(max(x_center, 0.0), 1.0)
        y_center = min(max(y_center, 0.0), 1.0)
        width = min(max(width, 0.0), 1.0)
        height = min(max(height, 0.0), 1.0)

        yolo_annotations.append((class_id, x_center, y_center, width, height))

    return yolo_annotations

def process_directory(json_dir, output_dir, class_name_to_id):
    os.makedirs(output_dir, exist_ok=True)

    for json_path in glob.glob(os.path.join(json_dir, "*.json")):
        yolo_annotations = convert_labelme_to_yolo(json_path, class_name_to_id)

        base_name = Path(json_path).stem
        txt_path = os.path.join(output_dir, f"{base_name}.txt")

        with open(txt_path, 'w') as f:
            for ann in yolo_annotations:
                f.write(f"{ann[0]} {ann[1]:.6f} {ann[2]:.6f} {ann[3]:.6f} {ann[4]:.6f}\n")
        print(f"Processed {json_path} -> {txt_path}")

if __name__ == "__main__":
    # 类别编号
    class_name_to_id = {
        "head": 0,
        "body": 1,
    }
    
    json_dir = "data/cs_label"
    output_dir = "data/cs_yolo_label"
    
    process_directory(json_dir, output_dir, class_name_to_id)
数据集分类

将数据集分为训练、验证和测试三个类别,一般来说占比大致为7:2:1,训练集和验证集是必须的,测试集可以不用,这边我们只有一百张图片,直接按照8:2的比例分为训练集和验证集。得到以下文件结构。

-data
	-train
		-images
		-labels
	-val
		-images
		-labels
	#  如果有测试集的话
	-test
		-images
		-labels

在data目录中新建一个data.yaml文件里面标记数据集位置以及分类信息

path: E:/Model/yolo11/data  # dataset root dir
train: train/images  # train images (relative to 'path')
val: val/images  # val images (relative to 'path')
test: test/images # test images (relative to 'path')

# Classes
nc: 2  # number of classes
names: ['head',  'body']  # class names

三、模型训练

​ 项目地址: https://github.com/ultralytics/ultralytics

​ 文档: https://docs.ultralytics.com/zh

训练

先在仓库中下载预训练的yolo检测模型,我们在yolo11n这个模型基础上进行训练。

请添加图片描述

写一段简短的代码即可开始训练

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO("yolo11n.pt")

    # Train the model
    results = model.train(data="data/data.yaml", epochs=300, imgsz=640, batch=4)

这边就讲几个常用的参数:

epochs:训练轮次,过高或过低都可能会影响训练效果,

imgsz:数据集图片尺寸大小

workers:加载数据的工作线程数(每 RANK 如果多GPU 训练)。影响数据预处理和输入模型的速度,尤其适用于多GPU 设置。

batch: 每次训练喂多少数据,有三种模式: 设置为整数(如 batch=16)、自动模式,内存利用率为 60%GPU (batch=-1),或指定利用率的自动模式 (batch=0.70).

执行前注意解释器要选择我们之前准备好的环境。

测试

训练好的模型放在runs文件夹下,用一段代码检测我们刚才训练的模型看检测结果是否准确。

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO("best.pt")
    
    # 结果验证
    results = model(r"data\train\images\CS130204.png")
    results[0].show()

请添加图片描述

结果还不错,头部的置信度稍微低一些,因为属于小目标检测。

导出onnx模型
from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO("best.pt")
    
   model.export(format='onnx')

四、结语

最后使用mss进行实时截图实战发现远处的物体识别率比较低,有可能是我置信度开的比较高的原因,还有就是速度有提升空间,每张照片识别基本在10ms这个样子加上一些后置的处理基本就消耗六七十ms了还是挺慢的,导致实时检测的图像帧率较低,后面优化后再看看控制鼠标移动轨迹,做个自瞄玩玩。

本文不涉及原理性的讲解,仅作为入门参考,带着训练一个自己的模型,主要是数据标注会比较浪费时间,如果是一些比较热门的项目,网上可能会有开源的数据集可以搜搜看,有问题可以评论区一起沟通交流。

Logo

尧米是由西云算力与CSDN联合运营的AI算力和模型开源社区品牌,为基于DaModel智算平台的AI应用企业和泛AI开发者提供技术交流与成果转化平台。

更多推荐