一文读懂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了还是挺慢的,导致实时检测的图像帧率较低,后面优化后再看看控制鼠标移动轨迹,做个自瞄玩玩。
本文不涉及原理性的讲解,仅作为入门参考,带着训练一个自己的模型,主要是数据标注会比较浪费时间,如果是一些比较热门的项目,网上可能会有开源的数据集可以搜搜看,有问题可以评论区一起沟通交流。