使用腾讯云高性能空间部署YOLOv11训练模型

背景

公司为传统工业互联网，内部基础零件较多，手动区分筛选费时费力，故而要采用视觉模型来自动标注与区分。本人Java开发，对于Python语言了解不多，且大模型基本上对性能要求很高，由于电脑配置本身，很难训练模型，所以采用腾讯高性能空间来学习。

环境部署

因为高性能空间本身部署了Python以及conda等基础服务，所以只需要配置yolo服务对应的组件即可。
1.PyTorch：-官方地址:https://pytorch.org/ 一定要选择和显卡型号对应的CUDA版本，没有英伟达的显卡，可以选择安装CPU版本。
2.labelImg:模型标注工具，可直接pip安装。
3.项目源码：https://github.com/ultralytics/ultralytics，解压源码后再按照要求补充环境依赖。我是直接在项目本地根目录下创建了一个requirements.txt依赖文件。

# Ultralytics requirements
# Example: pip install -r requirements.txt

# Base ----------------------------------------
matplotlib>=3.3.0
numpy==1.24.4 # pinned by Snyk to avoid a vulnerability
opencv-python>=4.6.0
pillow>=7.1.2
pyyaml>=5.3.1
requests>=2.23.0
scipy>=1.4.1
tqdm>=4.64.0

# Logging -------------------------------------
# tensorboard>=2.13.0
# dvclive>=2.12.0
# clearml
# comet

# Plotting ------------------------------------
pandas>=1.1.4
seaborn>=0.11.0

# Export --------------------------------------
# coremltools>=7.0  # CoreML export
# onnx>=1.12.0  # ONNX export
# onnxsim>=0.4.1  # ONNX simplifier
# nvidia-pyindex  # TensorRT export
# nvidia-tensorrt  # TensorRT export
# scikit-learn==0.19.2  # CoreML quantization
# tensorflow>=2.4.1  # TF exports (-cpu, -aarch64, -macos)
# tflite-support
# tensorflowjs>=3.9.0  # TF.js export
# openvino-dev>=2023.0  # OpenVINO export

# Extras --------------------------------------
psutil  # system utilization
py-cpuinfo  # display CPU info
thop>=0.1.1  # FLOPs computation
# ipython  # interactive notebook
# albumentations>=1.0.3  # training augmentations
# pycocotools>=2.0.6  # COCO mAP
# roboflow

数据准备

1.在安装好labelimg后，使用命令（labelImg）将其打开，进行数据标注，可参考：https://blog.csdn.net/ChaoChao66666/article/details/129129621 。数据标注好之后，一定要按照图片和标签区分开，并文件名称除后缀外其他都一致。在项目目录下创建训练基础数据和标注的文件夹，我创建的是data/images和data/labels，后面需要使用。
2.数据切分：在项目根目录下创建Python脚本文件，我定义的是main.py。如下：

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os

# 原始路径
image_original_path = "data/images/"
label_original_path = "data/labels/"

cur_path = os.getcwd()
# 训练集路径
train_image_path = os.path.join(cur_path, "datasets/images/train/")
train_label_path = os.path.join(cur_path, "datasets/labels/train/")

# 验证集路径
val_image_path = os.path.join(cur_path, "datasets/images/val/")
val_label_path = os.path.join(cur_path, "datasets/labels/val/")

# 测试集路径
test_image_path = os.path.join(cur_path, "datasets/images/test/")
test_label_path = os.path.join(cur_path, "datasets/labels/test/")

# 训练集目录
list_train = os.path.join(cur_path, "datasets/train.txt")
list_val = os.path.join(cur_path, "datasets/val.txt")
list_test = os.path.join(cur_path, "datasets/test.txt")

train_percent = 0.8
val_percent = 0.1
test_percent = 0.1


def del_file(path):
    for i in os.listdir(path):
        file_data = path + "\\" + i
        os.remove(file_data)


def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    else:
        del_file(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)
    else:
        del_file(train_label_path)

    if not os.path.exists(val_image_path):
        os.makedirs(val_image_path)
    else:
        del_file(val_image_path)
    if not os.path.exists(val_label_path):
        os.makedirs(val_label_path)
    else:
        del_file(val_label_path)

    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    else:
        del_file(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)
    else:
        del_file(test_label_path)


def clearfile():
    if os.path.exists(list_train):
        os.remove(list_train)
    if os.path.exists(list_val):
        os.remove(list_val)
    if os.path.exists(list_test):
        os.remove(list_test)


def main():
    mkdir()
    clearfile()

    file_train = open(list_train, 'w')
    file_val = open(list_val, 'w')
    file_test = open(list_test, 'w')

    total_txt = os.listdir(label_original_path)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)

    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
    num_test = num_txt - num_train - num_val

    train = random.sample(list_all_txt, num_train)
    # train从list_all_txt取出num_train个元素
    # 所以list_all_txt列表只剩下了这些元素
    val_test = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    val = random.sample(val_test, num_val)

    print("训练集数目：{}, 验证集数目：{}, 测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]

        srcImage = image_original_path + name + '.jpg'
        srcLabel = label_original_path + name + ".txt"

        if i in train:
            dst_train_Image = train_image_path + name + '.jpg'
            dst_train_Label = train_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
            file_train.write(dst_train_Image + '\n')
        elif i in val:
            dst_val_Image = val_image_path + name + '.jpg'
            dst_val_Label = val_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
            file_val.write(dst_val_Image + '\n')
        else:
            dst_test_Image = test_image_path + name + '.jpg'
            dst_test_Label = test_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_test_Image)
            shutil.copyfile(srcLabel, dst_test_Label)
            file_test.write(dst_test_Image + '\n')

    file_train.close()
    file_val.close()
    file_test.close()


if __name__ == "__main__":
    main()

执行此脚本后，会在文件目录下，自动创建好datasets相关内容，具体目录功能和代码中对应，格式如下：

模型训练

在项目目录下创建train.py，并粘贴如下代码：

from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(r'yolov11.yaml')  # 此处以 m 为例，只需写yolov11m即可定位到m模型
    model.train(data=r'data.yaml',
                imgsz=640,
                epochs=10, #轮次
                single_cls=True,  # 多类别设置False
                batch=16,
                workers=10,
                device='0', #如果没有显卡，这里写cpu
                )

可以看到，文件中还应用了其他文件，所以同样的，要创建data.yaml。

path: ../datasets  # 数据集所在路径
train: train.txt  # 数据集路径下的train.txt
val: val.txt  # 数据集路径下的val.txt
test: test.txt  # 数据集路径下的test.txt

# Classes
names:
  0: wave

model = YOLO(r’yolov11.yaml’)这里可以使用源码中的，也可以自己复制一份放在目录下，源码位置在：ultralytics/cfg/models，执行train.py，会生成对应的tran文件，里面包含权重文件，.pt后缀，一般生成路径位置是根目录下：/runs/detect/train

模型验证

编写验证脚本，val.py

import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('runs/detect/train/weights/best.pt') #上面生成的权重文件，我用的是文件路径方式
    model.val(data='data.yaml',
              imgsz=640,
              batch=16,
              split='test',
              workers=10,
              device='0',
              )

执行之后，可以看到生成的验证报告，路径：runs/detect/val

预测

编写脚本：detect.py，使用生成的权重文件，进行测试验证：

import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('runs/detect/train/weights/best.pt')
    model.predict(source='E://doc//yolo//test//images//23.jpg',  #本地测试文件地址
                  imgsz=640,
                  device='0',
                  save=True
                  )

执行之后，会看到标注这个文件已经被标注识别了，文件路径在runs/detect/predict。

总结

比较简单容易上手，也可以使用官方文档去学习参考，具体操作需要查看对应的源码和资料学习，以上仅为个人学习经验，不足之处还望理解，谢谢~