计算机视觉检测之yolo模型训练

基于yolo11n.pt模型，在新数据集上训练，使模型能够识别新事物

yolo11n.pt模型识别汽车结果如下：

在新数据集上训练，得到新模型shu.pt，新模型能够识别如下图片：

一、环境准备

安装minconda

创建虚拟环境 ：conda create -n yolo python=3.12

进入虚拟环境：conda activate yolo

在虚拟环境安装以下第三方库

 pip install torch torchvision torchaudio jupyter

前三个必须安装，最后一个可以不用安装。这里torch torchvision torchaudio都是cpu版本，因为数据和模型都不大，本次训练可以在cpu进行。有cuda的电脑可以安装cuda版本的，训练速度会更快。

二、程序准备

下载ultralytics-8.3.55源代码，需要把它编译成python的第三方库

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

此处下载的源码版本是ultralytics-8.3.55

下载后解压，这里使用pycharm打开，选用上面的虚拟环境yolo作为编译环境

**切换目录到ultralytics-8.3.55（源代码所在的文件夹），**命令行输入以下指令，将ultralytics-8.3.55安装到虚拟环境中，作为第三方库

pip install -e .

然后命令行敲入pip list，看到第三方库已经生成

三、数据准备

使用labelImg.exe标注训练数据

地址：https://github.com/HumanSignal/labelImg/releases/tag/v1.8.1

数据标注

这里一共准备了51张图片，其中10张作为验证集，40张作训练集，数据放置目录如下：

四、训练

模型准备，下载yolo11n.pt模型

地址：https://docs.ultralytics.com/zh/models/yolo11/#performance-metrics 需要科学上网

将yolo11n.pt模型 放在 ultralytics-8.3.55 文件夹下。
datasets、icon.yaml、train.py, yolo11n.pt的放置位置如下：

icon.yaml的内容为：

path: icon # dataset root dir
train: images/train
val: images/val
test:  

# Classes
names:
  0: shu

新建一个train.py文件训练模型，内容如下：

from ultralytics import YOLO

# Load a model
model = YOLO("yolo11n.pt")  # load a pretrained model (recommended for training)

# Train the model
results = model.train(data="icon.yaml", workers = 0, epochs=300, batch=16)

点击运行即可训练模型。

五、测试

本次训练只在cpu上进行，大概耗时1小时50分。训练后，模型所在位置：

将模型复制到与train.py所在目录同级的位置，改名为shu.pt

现在，可以使用新训练的模型预测数据啦（source是待预测的图片地址）：

from ultralytics import YOLO
yolo = YOLO(model='./shu.pt', task='detect')
result = yolo(source=r'E:\code\ultralytics-8.3.55\datasets\icon\test\1350.jpg', save = True)

上述预测代码执行后，结果会保存在runs/detect/predict文件夹下

打开即可。(左图是原图 右图是模型检测后的图片)

原文链接：https://www.bilibili.com/video/BV18EWyzHE4j