Datawhale AI 从零上手CV竞赛Task2：YoloV8训练视觉模型代码的优化

1.源代码

原本代码的内容如下:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import warnings
warnings.filterwarnings('ignore')


from ultralytics import YOLO
model = YOLO("yolov8m.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=32, imgsz=1080, batch=16)

2.上述代码分析

上述使用了YOLOv8模型进行训练，代码中包含以下关键步骤：

1、设置CUDA设备：

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

这行代码指定了使用哪块GPU进行训练。"0"表示使用第一块GPU。如果有多块GPU，你可以指定其他的编号。

要查看当前系统中有多少个可用的GPU，你可以使用Python中的torch库。以下是使用torch检查GPU数量的方法：

import torch

# 检查是否有可用的GPU
if torch.cuda.is_available():
    print("CUDA is available!")
    # 获取GPU的数量
    num_gpus = torch.cuda.device_count()
    print(f"Number of GPUs available: {num_gpus}")
else:
    print("CUDA is not available.")

这段代码首先检查CUDA是否可用，如果可用，则打印出可用的GPU数量。

2、忽略警告：

warnings.filterwarnings('ignore')

这行代码会忽略所有的警告信息，确保输出的日志更简洁。

3、加载YOLOv8模型：

model = YOLO("yolov8m.pt")

使用了yolov8m.pt模型，这是YOLOv8的中等版本（medium）。YOLOv8有多个版本，yolov8n.pt (nano)，yolov8s.pt (small)，yolov8m.pt (medium)，yolov8l.pt (large)，yolov8x.pt (extra-large)。每个版本的参数量和计算需求不同。

4、训练模型：

results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16)

在这个代码中设置了训练的参数：

• data="yolo-dataset/yolo.yaml"：定义了数据集的配置文件。
• epochs=64：训练64个周期。
• imgsz=1088：训练时使用的图像尺寸为1088x1088像素。
• batch=16：每个训练批次的大小为16张图片。

3.提高精度的建议及对应代码

1、增加数据增强

数据增强是提高模型泛化能力的有效方法。你可以在yolo.yaml配置文件中添加或增强数据增强技术，如随机旋转、裁剪、缩放等。

# yolo.yaml 文件中的部分内容
augment:
  flipud: 0.5    # 随机上下翻转
  fliplr: 0.5    # 随机左右翻转
  hsv_h: 0.015   # HSV色调调整
  hsv_s: 0.7     # HSV饱和度调整
  hsv_v: 0.4     # HSV亮度调整
  translate: 0.1 # 平移
  scale: 0.5     # 缩放
  shear: 0.01    # 剪切
  perspective: 0.001 # 透视变换

2、调用多个GPU

如果你想使用所有可用的GPU，可以设置CUDA_VISIBLE_DEVICES环境变量为空或者包含所有GPU的编号。例如，如果你有4个GPU，可以设置为0,1,2,3。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"  # 假设有4个GPU

3、使用更大的模型

如果你的GPU允许，可以考虑使用更大的模型（例如yolov8x.pt）来提高精度。这会增加模型的容量，但也需要更多的计算资源。（注：我在使用4090GPU24G显存的情况下，跑yolov8m模型进行微调大约用了19.8G的显存，而yolov8l与yolov8x则需要远远高于m的显存，小心爆显存！）

model = YOLO("yolov8x.pt")

4、调整学习率调度器

学习率调度器可以帮助模型在训练过程中更好地收敛。YOLOv8内部已经包含了自动调整学习率的机制，但你可以手动设置lr0（初始学习率）和lrf（最终学习率）来微调。

results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, lr0=0.001, lrf=0.01)

5、增加正则化

可以通过设置weight_decay来增加正则化，防止模型过拟合。

results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, weight_decay=0.0005)

6、使用提前停止

提前停止是一种防止过拟合的方法，当模型在验证集上不再改善时提前停止训练。

results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, patience=10)

4.综合改进后的代码

综合以上建议，改进后的代码可能如下：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

import warnings
warnings.filterwarnings('ignore')

from ultralytics import YOLO

# 使用更大的模型
model = YOLO("yolov8x.pt")

# 训练模型，包含数据增强、学习率调整、正则化和提前停止
results = model.train(
    data="yolo-dataset/yolo.yaml", 
    epochs=32, 
    imgsz=1088, 
    batch=16, 
    lr0=0.001, 
    lrf=0.01, 
    weight_decay=0.0005, 
    patience=10
)