[卷积神经网络]使用YOLOv11训练自己的模型

YoloV11的源码：YOLOv11

一、数据集准备

首先，准备好自己的数据集，包含图像文件和标注文件，因为我的数据集上Voc格式，所以需要先转为yolo格式，可以使用下面的脚本进行转换。

import os
import shutil
import xml.etree.ElementTree as ET# 类别映射 (需要记住这个顺序，后续YOLO的配置文件中也要与这个顺序一致)
class_mapping = {'cls1':0,'cls2':1,# ...
}def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):tree = ET.parse(voc_annotation_file)root = tree.getroot()size = root.find('size')width = float(size.find('width').text)height = float(size.find('height').text)with open(yolo_label_file, 'w') as f:for obj in root.findall('object'):cls = obj.find('name').textif cls not in class_mapping:continuecls_id = class_mapping[cls]xmlbox = obj.find('bndbox')xmin = float(xmlbox.find('xmin').text)ymin = float(xmlbox.find('ymin').text)xmax = float(xmlbox.find('xmax').text)ymax = float(xmlbox.find('ymax').text)x_center = (xmin + xmax) / 2.0 / widthy_center = (ymin + ymax) / 2.0 / heightw = (xmax - xmin) / widthh = (ymax - ymin) / heightf.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")if __name__ == '__main__':# VOC格式数据集路径voc_data_path = 'VocData'    #请修改为自己的Voc数据集路径voc_annotations_path = os.path.join(voc_data_path, 'Annotations')voc_images_path = os.path.join(voc_data_path, 'JPEGImages')# YOLO格式数据集保存路径yolo_data_path = 'datasets/VOC'    #请修改为自己的输出路径yolo_images_path = os.path.join(yolo_data_path, 'images')yolo_labels_path = os.path.join(yolo_data_path, 'labels')# 创建YOLO格式数据集目录os.makedirs(yolo_images_path, exist_ok=True)os.makedirs(yolo_labels_path, exist_ok=True)print('开始转换')i = 0# 遍历VOC数据集的Annotations目录，进行转换for voc_annotation in os.listdir(voc_annotations_path):if voc_annotation.endswith('.xml'):voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)image_id = os.path.splitext(voc_annotation)[0]voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg")yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg")convert_voc_to_yolo(voc_annotation_file, yolo_label_file)if os.path.exists(voc_image_file):shutil.copy(voc_image_file, yolo_image_file)i+=1print('{}/{}'.format(i,len(os.listdir(voc_annotations_path))))print("转换完成！")

在完成数据集格式的转换后还要对其进行划分，使用下面的代码将数据集划分为train,val,test

import os, shutil
from sklearn.model_selection import train_test_splitval_test_size = 0.2    #验证集占比
test_size = 0.5
postfix = 'jpg'
imgpath = 'datasets/VOC/images'    #数据集图像文件夹
txtpath =  'datasets/VOC/labels'   #数据集标签文件夹output_folder = 'datasets\TestA'    #输出文件夹output_train_img_folder = output_folder +'/images/train'
output_val_img_folder =  output_folder +'/images/val'
output_test_img_folder =  output_folder +'/images/test'
output_train_txt_folder =  output_folder +'/labels/train'
output_val_txt_folder =  output_folder +'/labels/val'
output_test_txt_folder =  output_folder +'/labels/test'os.makedirs(output_train_img_folder, exist_ok=True)
os.makedirs(output_val_img_folder, exist_ok=True)
os.makedirs(output_test_img_folder, exist_ok=True)
os.makedirs(output_train_txt_folder, exist_ok=True)
os.makedirs(output_val_txt_folder, exist_ok=True)
os.makedirs(output_test_txt_folder, exist_ok=True)listdir = [i for i in os.listdir(txtpath) if 'txt' in i]
train, val_test = train_test_split(listdir, test_size=val_test_size, shuffle=True, random_state=0)test, val = train_test_split(val_test, test_size=test_size, shuffle=True, random_state=0)idx = 0for i in train:img_source_path = os.path.join(imgpath, '{}.{}'.format(i[:-4], postfix))txt_source_path = os.path.join(txtpath, i)img_destination_path = os.path.join(output_train_img_folder, '{}.{}'.format(i[:-4], postfix))txt_destination_path = os.path.join(output_train_txt_folder, i)shutil.copy(img_source_path, img_destination_path)shutil.copy(txt_source_path, txt_destination_path)idx += 1print('{}/{}'.format(idx,len(train)+len(val)+len(test)))for i in val:img_source_path = os.path.join(imgpath, '{}.{}'.format(i[:-4], postfix))txt_source_path = os.path.join(txtpath, i)img_destination_path = os.path.join(output_val_img_folder, '{}.{}'.format(i[:-4], postfix))txt_destination_path = os.path.join(output_val_txt_folder, i)shutil.copy(img_source_path, img_destination_path)shutil.copy(txt_source_path, txt_destination_path)idx+=1print('{}/{}'.format(idx,len(train)+len(val)+len(test)))for i in test:img_source_path = os.path.join(imgpath, '{}.{}'.format(i[:-4], postfix))txt_source_path = os.path.join(txtpath, i)img_destination_path = os.path.join(output_test_img_folder, '{}.{}'.format(i[:-4], postfix))txt_destination_path = os.path.join(output_test_txt_folder, i)shutil.copy(img_source_path, img_destination_path)shutil.copy(txt_source_path, txt_destination_path)idx+=1print('{}/{}'.format(idx,len(train)+len(val)+len(test)))print('结束分割')

二、配置文件

接下来，需要为YOLO配置如下文件，以确保其能正常工作。

1.配置数据集

数据集的映射关系如下图：

为其配置一个yaml文件，包含数据集的路径，类别信息。

# 数据集的配置yaml文件# 这里请写绝对路径，分别填入三个数据集中图像所在的位置
train: \yolov11\datasets\VOC\images\train  
val: \yolov11\datasets\VOC\images\val  
test: \yolov11\datasets\VOC\images\test  # 类别数
nc: 6# 类别名，请与之前转换时的序列保持一致
names: ['Boerner','Leconte','acuminatus','armandi','coleoptera','Linnaeus']

2.配置模型

新建一个yaml文件，你可以放在任何地方，这里为了方便管理放在了datasets目录下。且这个文件是官方提供的，不需要做过多更改：

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect# Parameters
nc: 6 # 类别数，与之前保持一致
scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'# [depth, width, max_channels]
#重点修改此位置，将你想要的那个模型配置放在第一位n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPsx: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPss: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPsm: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPsl: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs# YOLO11n backbone
backbone:# [from, repeats, module, args]- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4- [-1, 2, C3k2, [256, False, 0.25]]- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8- [-1, 2, C3k2, [512, False, 0.25]]- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16- [-1, 2, C3k2, [512, True]]- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32- [-1, 2, C3k2, [1024, True]]- [-1, 1, SPPF, [1024, 5]] # 9- [-1, 2, C2PSA, [1024]] # 10# YOLO11n head
head:- [-1, 1, nn.Upsample, [None, 2, "nearest"]]- [[-1, 6], 1, Concat, [1]] # cat backbone P4- [-1, 2, C3k2, [512, False]] # 13- [-1, 1, nn.Upsample, [None, 2, "nearest"]]- [[-1, 4], 1, Concat, [1]] # cat backbone P3- [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)- [-1, 1, Conv, [256, 3, 2]]- [[-1, 13], 1, Concat, [1]] # cat head P4- [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)- [-1, 1, Conv, [512, 3, 2]]- [[-1, 10], 1, Concat, [1]] # cat head P5- [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)- [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)

三、开始训练

编写一个训练脚本，然后直接运行它。

from ultralytics import YOLOif __name__ == '__main__':# 加载模型#model = YOLO('datasets/yolo11s.pt')    #预训练权重model = YOLO(model='datasets/yolo11.yaml')# 训练模型train_results = model.train(data = 'datasets/VOC/VOC.yaml',   #数据集路径epochs = 100,        #训练轮次imgsz = 640,         #图片尺寸device = '0',        #使用设备batch = 16,name = 'Result',    #项目保存路径)