Datawhale AI 夏令营 Task3：上分思路——数据集增强与模型预测

在前面两次Task中, 我们分别跑通了baseline, 并对baseline的epochs和batch_size进行了适当调整, 取得了比baseline高很多的得分.

在本次Task3中, 我们将对数据进行预处理, 同时对模型其他参数进行调整, 争取获得更好的检测效果.

1.扩充训练集和验证集

扩充训练和验证数据集, 能够让模型学到更多的特征, 提高了模型检测违规行为的能力. 此处我选取了训练集中的前20个视频和验证集中的后5个视频作为数据集, 如果想要进一步提高模型预测效果, 可以将所有的训练集和测试集选中, 此外还可以自己额外增加训练集, 因为YOLO的图像检测属于监督学习, 所以还需要增加对应的带有标签(标注)的验证集. 核心代码如下:

for anno_path, video_path in zip(train_annos[:20], train_videos[:20]): # 训练集
...
for anno_path, video_path in zip(train_annos[-5:], train_videos[-5:]): # 验证集
...

2.训练集预处理

YOLO模型对图像增强的参数如下:

考虑到视频中的图像存在白天和夜晚场景, 白天光线充足容易辨认, 但是夜晚各种灯光照射, 原有的目标可能相较于白天发生了很大的变化, 所以我们可以进行如下设置, 增强模型在夜晚图像中目标检测的抗干扰能力.

hsv_h = 0.2, hsv_s = 0.7, hsv_v = 0.4, bgr = 0.1

由于视频监控中的车辆可能以任意角度出现在画面中, 我们可以增强模型对目标旋转之后的检测, 除此之外, 测量不断移动也会造成目标在画面中的大小发生变化, 于是我们加上对目标缩放的之后的检测, 模拟目标与相机处于不同距离的情景.

 degrees = 20, scale = 0.5, flipud = 0.2, fliplr = 0.2,

目标还可能被其他干扰物遮挡, 于是我们加上遮挡和消除方面的数据增强:

copy_paste = 0.2, erasing = 0.1

参考代码如下:

import os
import warnings
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
warnings.filterwarnings('ignore')
from ultralytics import YOLO
epochs = 10
batch_size = 32
model = YOLO("yolov8l.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=epochs, imgsz=720, batch=batch_size, 
hsv_h = 0.02, hsv_s = 0.7, hsv_v = 0.4, degrees = 20, scale = 0.5, flipud = 0.2, fliplr = 0.2,
bgr = 0.1, mixup = 0.1, copy_paste = 0.2, erasing = 0.1)

3.减少图片像素, 加快训练速度

由于训练集都是视频监控, 被检测的目标其实只占整张图片的一小部分, 我们可以适当降低imgsz, 减小单张图片的显存占用, 同时可以增大batch_size, 让模型一次性能学到目标检测物更多的特征, 而且能保证训练时loss波动不会太大.

参数imgsz可以传入一个数值, 如imgsz = 1080, 此时模型会自动切分出1080*1080像素的图片进行目标检测训练. 而原视频分辨率是1920*1080, 这样会导致大量无效信息被输入到模型中, 导致显存占用极高, 我们可以设置为imgsz = 720, 适当减少信息输入, 减轻显存压力.
此外, 如果显存和算力都很富余, 也可以传入原始尺寸的图像, 即imgsz = [1920, 1080].

经过实测, 在上述参数下, 显存几乎能恰好贴着24G高压线跑而且不会爆显存, 这样可以最大化利用4090的算力, 加快训练速度.

训练的效果如下:

4.使用最佳权重进行测试

训练完成之后, 我们需要选择训练好的权重进行测试并输出result. 这里需要注意的是, 每重新训练一次都会在runs/detect/下生成一个新的train文件夹, 最后面带有2, 3…等编号, 每次训练完之后, 需要手动更改best.pt的路径, 避免拿以前训练的权重进行测试.

model = YOLO("runs/detect/train/weights/best.pt") # 这里需要改成最新的训练权重

当然我们可以增加如下代码自动选择最新的train文件夹下的best.pt.

from ultralytics import YOLO
import os
import re

def get_max_suffix_number(directory):
    max_number = 0
    pattern = re.compile(r'train(\d+)')

    for folder in os.listdir(directory):
        match = pattern.match(folder)
        if match:
            number = int(match.group(1))
            if number > max_number:
                max_number = number
    return max_number

directory_path = 'runs/detect'
max_number = get_max_suffix_number(directory_path)
model = YOLO("runs/detect/train/weights/best.pt") # 一定要保证train/weights/下至少有一个best.pt, 否则会报错
if max_number > 1:
    model = YOLO(f"runs/detect/train{max_number}/weights/best.pt", max_number)

5.最终得分

使用上述方法训练得到最终权重best.pt, 经过测试之后我们得出result文件, 打包并上传至比赛官网, 最终得分0.19

6.数据分析与思考

各种loss曲线

result.csv

通过分析YOLO模型自带保存的loss曲线, 我们发现val/box_loss和val/cls_loss居然有所上升, 可能原因是模型出现了过拟合. 此外, 还可能是因为默认学习率太高导致模型权重更新不稳定等. 如果时间充足, 我们还可以将epochs适当增大, 可能增大到一定轮数之后, 上述loss可能会下降并趋于稳定.