Datawhale AI 夏令营 Task2：建模方案解读与进阶

Task2：建模方案解读与进阶

在baseline的基础上(链接), 我们想要获得更好的检测效果和更高的得分, 需要在baseline代码上做出适当的调整.

1.更改训练集大小

找到baseline中如下代码:

for anno_path, video_path in zip(train_annos[:5], train_videos[:5]): # 训练集大小 建议改成10, 过大会导致训练集太多, 导入耗时
    print(video_path)
    anno_df = pd.read_json(anno_path)
    cap = cv2.VideoCapture(video_path)
    frame_idx = 0 
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        img_height, img_width = frame.shape[:2]
        
        frame_anno = anno_df[anno_df['frame_id'] == frame_idx]
        cv2.imwrite('./yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)

        if len(frame_anno) != 0:
            with open('./yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.txt', 'w') as up:
                for category, bbox in zip(frame_anno['category'].values, frame_anno['bbox'].values):
                    category_idx = category_labels.index(category)
                    
                    x_min, y_min, x_max, y_max = bbox
                    x_center = (x_min + x_max) / 2 / img_width
                    y_center = (y_min + y_max) / 2 / img_height
                    width = (x_max - x_min) / img_width
                    height = (y_max - y_min) / img_height

                    if x_center > 1:
                        print(bbox)
                    up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
        
        frame_idx += 1

这段代码的含义是:选取列表train_anno, train_videos中的前5个元素作为训练集, 并以.jpg和.txt格式保存到./yolo-dataset/train/目录下, 如果我们将train_annos[:5], train_videos[:5]中的5改成10, 可以得到更多的训练数据, 模型可以学习到更多的特征. 但是参数太大会导致占用大量硬盘空间, 而且导入数据也会消耗大量时间, 建议多次调整追求训练效果和训练耗时的最佳平衡.

2.改用更好的模型

baseline中使用的是YOLOv8n模型, 是v8中最基础的版本, 其优点是训练速度快, 算力要求低, 可以在资源有限的情况下进行目标检测, 除此之外还有v8s/v8m/v8l/v8x, 具体参数如下:

Model	size	mAP	Speed(ms)	params(M)	FLOPs(B)
YOLOv8n(点击下载模型)	640	37.3	0.99	3.2	8.7
YOLOv8s	640	44.9	1.20	11.2	28.6
YOLOv8m	640	50.2	1.83	25.9	78.9
YOLOv8l	640	52.9	2.39	43.7	165.2
YOLOv8x	640	53.9	3.53	68.2	257.8

我们选择更换为v8s, 将代码

!wget http://mirror.coggle.club/yolo/yolov8n-v8.2.0.pt -O yolov8n.pt

更改为

!wget http://mirror.coggle.club/yolo/yolov8s-v8.2.0.pt -O yolov8s.pt

并将调用模型的代码

from ultralytics import YOLO
model = YOLO("yolov8n.pt")

更改为

model = YOLO("yolov8s.pt")

最新版的YOLOv10已经发布, 想要冲击更高的分的可以点这前往YOLOv10主页
注意:使用参数更大的模型可能导致占用更多的显存, 如果代码运行报错为显存不足, 可以适当减小batch

3.调整训练参数

baseline中

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

训练参数中的epochs表示在训练集上经过完整训练的轮数, imgsz是image size的缩写, 表示图像分辨率, imgsz = 1080表示图像分辨率是1920*1080(16:9), batch表示一次训练同时使用的样本数
适当增大epochs, 可以进一步降低训练loss, 得到更为精确拟合训练集的模型, 从而提高检测精准度. 但是过大的epochs可能导致过拟合, 即模型在训练集上可以很好地检测出目标, 但是在其他数据集(如测试集, 验证集)中表现远远弱于训练集, 即对训练集形成了一种过度依赖关系.
适当增大batch可以使得模型在同一批次中学习到更多的特征, 还能减少训练批次, 降低时间成本, 缺点是需要更大的内存和显存以便同时存放同一批次的训练数据.
我们将其改成

results = model.train(data="yolo-dataset/yolo.yaml", epochs=10, imgsz=1080, batch=32)

最后我们将调整好的代码上传到服务器上, 一键运行, 得到result.zip后提交到比赛官网上, 得分为0.065, 比第一次得分0.0047高了很多.

此次只是粗略地调整参数, 为了得到更高的分数, 后续可能需要进一步调整参数和数据增强, 如对视频中抽帧得到的图像进行反转旋转, 更改对比度饱和度等.