使用YOLO v10x训练BDD100K数据集

一、背景

最近要参加一个比赛，比赛是做图像识别的，比赛提供了部分BDD100K数据集，每个图片是1280*720的交通图片，每个标签中带有物体类别、物体左上锚点、物体右下锚点，如下：

比赛提供的标签数据有9种目标，如下：

比赛前的培训课程上，官方演示了使用YOLO V5模型训练了部分数据，因此我准备使用更高版本的YOLO V10模型训练完整的BDD100K数据集
我的环境信息如下：

• 操作系统windows11
• 显卡为NVIDIA GeForce RTX 3080
• cuda版本11.7
• cudnn版本8.9.7.29
• torch版本2.0.1+cu117
• torchvision版本0.15.0+cu117

二、准备数据集与环境

1. 下载BDD100K数据集

我忘记我的数据集是从哪里下载的了，可以参照BDD100K官方文档去下载，不过我下载的数据集和官网上的路径不大一致，但是差别不大，可以用

我下载下来的数据集路径如下：

images是图片文件路径，里面分为10k和100k，10k的数据体量小，我这里要训练100k的数据，test、train、val分别是测试集、训练集和验证集，里面都是jpg格式的图片数据

labels是标签文件路径，我下载的这里只有100k的，是不是10k的标签在100k里都能找到我没有去验证，对应这images/100k里的路径，labels/100l里有训练集和验证集的标签文件，都是json格式的，文件名与图片一一对应，如下：


标签json文件格式如下：

需要注意的是，BDD100K的数据标签里objects列表后面有一些内容长这样：

我没有去深挖这个标签代表什么，有兴趣可以去研究，因为我在这个项目中用不到这个标签，因此后面做标签转换的时候略过了这些标签。

2. 下载YOLO V10代码

从YOLO V10的Github主页上下载代码，有git的用git，没有git直接下载zip即可

下载完成后，把YOLO v10的代码放到项目路径中，我的路径如图：

3. 标签转换

BDD100K的数据标注方式与YOLO v10的标注方式不同，BDD100K是用左上和右下两个锚点的标注方式，而YOLO V10则是是用的中心点标注方法，例如：

在上图这个例子里，图片中有三个物体，两个人和一条领带，yolo模型的标注文件是txt文件，一行一条标注信息

我们以图片上的齐达内为例，他的标注有5个信息，分别是类别、中心点x相对位置，中心点y相对位置，图像高百分比height，图像宽百分比height

yolo模型不管图片的像素尺寸是多少，图片的长和宽都是0到1的范围，里面的每个物体也是一个矩形，x=0.48表示齐达内的中心点在图像x轴上的位置为0.48，y=0.63表示齐达内的中心点在图像y轴的0.63，height=0.71表示齐达内的高占整个图片高的71%，width=0.69表示齐达内的宽占整个图片宽的69%，应该还是很好理解的。

这种标注相比于BDD100K的标注，好处在于更加通用，如果BDD100K的图像被缩放，不再是1280*720的尺寸，那么所有的标注锚点坐标都需要对应修改

所以我们需要先编写一个python脚本，来吧json格式的BDD100K标注文件转化成txt格式的YOLO v10标注文件，基本思路如下

假设我们有一张图片，大小是1280*720，在图片内部有一个标签为car的物体，左上角锚点的坐标是(100, 120)，右下角的锚点坐标是(200, 320)，那么：

• 中心点坐标X = ((200 - 100) / 2 + 100) / 1280 = 150 / 1280 = 0.1171875
• 中心点坐标Y = ((320 - 120) / 2 + 120) / 720 = 440 / 720 = 0.3055556
• 物体宽度width = (200 - 100) / 1280 = 0.078125
• 物体高度height = (320 - 120) / 720 = 0.277778
  

简化一下，得到公式：

• center_X = (x1 + x2) / 2 / 1280
• center_Y = (y1 + y2) / 2 / 720
• width = (x2 - x1) / 1280
• height = (y2 - y1) / 1280

根据上面的公示编写脚本就可以了，需要注意的是要把训练集的标签和验证集的标签都进行转换，我的转换脚本名为转换标签.py，路径在cs目录中：

代码如下：

import json
import os

# 输入和输出路径
# input_json_path = '../bdd100k/labels/100k/train'
# output_txt_dir = '../bdd100k/labels/100k/train'
input_json_path = '../bdd100k/labels/100k/val'
output_txt_dir = '../bdd100k/labels/100k/val'

# 类别映射，可以根据需要修改
category_map = {
    'bus': 0,
    'traffic light': 1,
    'traffic sign': 2,
    'person': 3,
    'bike': 4,
    'truck': 5,
    'motor': 6,
    'car': 7,
    'rider': 8,
    # 添加其他类别
}

category_keys_list = list(category_map.keys())  # 所有需要识别的物体名称列表

# 确保输出目录存在
os.makedirs(output_txt_dir, exist_ok=True)


def change_label_type(filename):
    """
    修改一个文件的标签类别
    :param filename: json路径
    :return:
    """
    with open(filename) as f:
        data = json.load(f)
    file_name = data.get("name")
    frames = data.get("frames")
    with open(f"{output_txt_dir}/{file_name}.txt", "w", encoding="utf-8") as out_file:
        for each_frame in frames:
            objects = each_frame.get("objects")
            for each_object in objects:
                category = each_object.get("category")
                if category not in category_keys_list:
                    continue
                bbox = each_object['box2d']
                category_id = category_map.get(category)
                xmin, ymin, xmax, ymax = bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2']
                x_center = (xmin + xmax) / 2 / 1280
                y_center = (ymin + ymax) / 2 / 720
                width = (xmax - xmin) / 1280
                height = (ymax - ymin) / 720

                out_file.write(f"{category_id} {x_center} {y_center} {width} {height}\n")


# 读取BDD100K的标注文件
bdd100k_labels_list = os.listdir(input_json_path)
for label_file in bdd100k_labels_list:
    label_file_path = f"{input_json_path}/{label_file}"
    change_label_type(label_file_path)

上面的脚本要执行两次，一次转换训练集的标签，一次转换验证集的标签，方法是依次注释或解除注释input_json_path和output_txt_dir这几行代码，input_json_path表示bdd100k的标签路径，output_txt_dir表示转换后的yolo标签输出路径，这里输出路径就直接输出到原标签路径，也就是转换前的json标签和转换后的txt标签在在同一个路径中

脚本中的category_map是参照比赛说明文件编写的（在本文最开头有一个图片表格），因为bdd100k里面用文本而不是数字做为物体类别的标注，为了和比赛一致，所以做了物体名称和数字编号的映射

得到的标签文件如图：

yolo模型的标签文件是txt格式，每行标注一个物体，每行有5个信息，用空格分割，分别是物体类别编号、中心点x坐标、中心点y坐标，物体宽、物体高

4. 开发环境准备

参照YOLO v10 Github主页上的说明，用conda创建一个名为yolov10的环境，执行下面的代码：
conda create -n yolov10 python=3.9
期间需要输入一个y然后按回车确认
安装完成后激活环境
conda activate yolov10
随后安装所需依赖，依赖列表文件requirements.txt在yolo v10的源码中，所以我们需要先进入到源码路径
这个步骤就不贴命令了，每个人路径都不一定一样，直接上图

requirements.txt文件就在yolov10-main目录中

执行命令pip install -r requirements.txt安装依赖，这个过程比较耗时，耐心等待，如果网络不佳，可以尝试切换conda镜像源，方法不在这里赘述，网上很多教程

然后执行pip install -e .安装yolo命令集，完成开发环境准备

5. GPU环境确认

这里需要注意一点，YOLO v10提供的requiremetns.txt中，torch的版本是2.0.1，这就需要torch版本和CUDA、cuDNN版本对应起来

在torch下载的网站上可以看到，torch-2.0.1版本对应的cuda版本为11.7

上图cu117表示cuda版本为11.7，cp39表示python版本为3.9，win表示windows系统

所以如果你的CUDA不是11.7版本，很有可能无法使用GPU训练，可以使用nvcc --version

显然我的cuda版本是11.2，使用这个版本到了后面也确实无法使用GPU训练，因为torch检测不到GPU设备

FUCK NVIDIA！！

安装CUDA和CUDNN，参照我的另一篇文章，这里不再赘述过程了

还有一点需要注意，那就是YOLO v10的源码包中的requirements.txt里面的torch是cpu版本的，如果要使用GPU，需要去下载GPU版本的torch以及对应版本的torchvision，保险起见仍然使用torch==2.0.1版本，但是需要找到GPU版本

在torch下载的网站上找到的带有cu开头的都是GPU版本，因为我是python3.9环境，CUDA版本为了适配torch=2.0.1升级到了11.7，因此我下载标红的那一个

然后从github主页上查看torch和torchvision的版本对应

因为torch版本是2.0.1，属于2.0系列，所以torchvision的版本需要是0.15，还是在torch下载的网站上找到对应的torchvision版本

安装的时候使用pip install xxxx.whl安装即可，注意要先安装torchvision，然后安装torch，否则可能会给弄一个版本不对劲的torch！

最后在环境里运行下面这个小脚本，看看能否查看到GPU

import torch
import torchvision

print(torch.__version__)  # 打印torch版本
print(torchvision.__version__)  # 打印torchvision版本
print(torch.version.cuda)  # 打印cuda版本

# 查看是否有可用的GPU
if torch.cuda.is_available():
    # 获取GPU的数量
    num_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {num_gpus}")

    # 打印每个GPU的编号和名称
    for i in range(num_gpus):
        print(f"Device {i}: {torch.cuda.get_device_name(i)}")
else:
    print("No GPU available.")

我的输出结果如下：

2.0.1+cu117
0.15.0+cu117
11.7
Available GPUs: 1
Device 0: NVIDIA GeForce RTX 3080

说明torch检测到了我的显卡，我可以用GPU去训练了

有GPU一定用GPU，相比于CPU快的不是一点半点

6. YOLO v10样例

这一步可以不做

在YOLO v10源码路径有个app.py，通过这是一个简易的web应用，可以上传单个图片进行检测

在YOLO v10源码路径里，执行python app.py可以启动这个web应用，打开浏览器访问http://127.0.0.1:7860可以查看web应用的界面，如下：

左上方有个框，我们可以把图片拖进去，然后点击下方的Detect Objects按钮，等待一会儿就可以对图像中的物体进行识别了

三、使用YOLO v10模型

1. 下载权重文件

首先需要下载YOLO v10的模型，每个模型的下载地址如下，简单理解为越往后越厉害就可以-

• YOLOv10-N：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt
• YOLOv10-S：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt
• YOLOv10-M：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10m.pt
• YOLOv10-B：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10b.pt
• YOLOv10-L：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10l.pt
• YOLOv10-X：https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10x.pt

下载完成后把.pt文件放置在源码根目录即可，我下载的最后一个 YOLOv10-X

2. 编写yaml配置文件

在yolov10-main中新建一个名为yolov10x.yaml的文件，内容如下：

path: xxxxxx
train: images/100k/train
val: images/100k/val
test: images/100k/test

names:
  0: bus
  1: traffic light
  2: traffic sign
  3: person
  4: bike
  5: truck
  6: motor
  7: car
  8: rider

上面的yaml文件容易理解，path是数据集的路径，这里我写了绝对路径
train、val、test分别是相对于path的训练集路径、验证集路径以及测试集路径，train和val路径里应当既有图片数据也有标签数据
names是标签类别，我参照比赛要求的表格进行填写，一共九个类别

3. 训练模型

在yolov10-main路径中执行指令即可进行训练：
yolo task=detect mode=train data=yolov10x.yaml model=yolov10x.pt epochs=100 batch=2 device=0 plots=True workers=4

对于上述指令的简要说明如下：

• yolo：运行yolo程序
• task=detect:指定为检测任务
• model=train：指定为训练任务
• data=yolov10x.yaml：依照刚刚编写的yolov10x.yaml文件进行训练
• model=yolov10x.pt：指定下载的yolov10预训练权重文件
• epochs=100：设置训练轮次
• batch=2：设置训练集加载批次，主要是提高训练速度，具体得看你的显卡或者内存容量。如果显存大，则可以设置大一些
• device=0：指定训练设备，如果没有gpu，则令device=cpu，如果有一个gpu，则令device=0，有两个则device=0,1以此类推进行设置
• plots=True：指定在训练过程中生成图表（plots）。这可以帮助你可视化训练进度，如损失函数的变化等
• workers=4：设置进程数，如果显存大，可以往大了设置

运行指令后就开始训练了

估计会训练很久了，训练完了继续更新！

–训练完成后更新–

历时四天，训练了个稀巴烂

内存爆了报MemoryError报了两次，一次是在10/100 Epoch，一次是在40/50 Epoch，人崩溃了许久

每一个Epoch训练都需要接近一小时的时间，40/50这一次明明胜利在望了，看完樊振东vs张本智和的比赛我还过来看了一眼，还在训练，心满意足睡觉去了，没想到啊，看时间应该是刚躺下就报错了，早上起来人很崩溃

我内存有32G，看样子数据量还是太大，如何解决这个问题后面在讨论吧，我现在迫切的想知道训练完之后的步骤，于是又设了一次Epoch=2，只训练两步，历时两个小时顺利训练完成

可以看到训练好的模型有两个，一个是best.pt，一个是last.pt，可以理解成表现最好的模型和最后一步训练出来的模型，他们存放在YOLO v10源码路径的runs/detect/train5/weights目录中，train5是因为这是我第五次训练任务，然后我在train4目录中仍然可以找到best.pt和last.pt，因此得出一个好消息是虽然之前训练因为内存爆了中断，但是我仍然可以使用训练了一半的模型

整个训练完成后会利用10000张验证集图片来检测模型的各个性能指标：Instances, Box, R, mAP50, mAP50-95，对它们的解读是：

• Instances: 验证集中目标实例的数量。
• Box: 预测的边界框的准确性。
• R (Recall): 模型的召回率，表示模型能检测到的目标实例比例。
• mAP50: 在 IoU=0.50 时的平均准确率，反映模型的总体检测精度。
• mAP50-95: 综合 IoU=0.50 到 0.95 的平均准确率，是一个更严格的模型检测性能指标。

详细说明如下（ChatGPT提供）：

4. 评估模型

现在我有了train4和train5两次训练的模型，train4是在40/50内存爆了没训练完的模型，train5是只训练了两步的模型，只训练了两步的模型评估效果在上面展示了，各项指标都不理想，于是我准备看一下train4得到的模型效果如何

执行yolo val model=runs/detect/train4/weights/best.pt data=yolov10x.yaml batch=1 imgsz=640
得到的结果如下：

这个比赛任务只关注物体识别的数量，而不关注物体在图片的哪个位置，因此Box这一项不重要，R和mAP50和train5的模型相比都有所提升，但并不是太大，看样子还是得多训练才好

四、使用模型获取目标数量

比赛最终会提供一些新的图片，因此需要调用模型来识别新图片中的物体数量，并且转换成比赛所需的csv格式

1. 预测图片

假设我们最终得到了一个名为best.pt的模型，那么使用这个模型去预测图片可以在YOLO v10源码路径去执行yolo predict model=runs/detect/train4/weights/best.pt source=../../../bdd100k/images/100k/test/ imgsz=640 save_txt=True

这里source是待预测图片的路径，save_txt表示把结果用文本文件的形式进行存储

执行后可以在YOLO v10源码路径的runs/detect/predict文件夹中找到

跟train过程一样，predict后面的数字表示第几次预测任务，这样方便回顾之前的预测结果

我的在predict2文件夹中，文件夹里的.jpg图片是通过模型标注好的图片，例如：

然后labels文件夹里是文本标注的形式了，打开一个长这样

这个是跟我们前面转换成YOLO的标注方式一样的格式，这个比赛不注重物体位置，只注重物体类别，所以我们只关注每行的第一个数字即可，它对应了一种目标物体

2. 整理比赛所需的csv文件

比赛最终提交一个csv文件，文件有3列，分别表示图片名、人的数量和车的数量，其中人只对应行人，车有5中车

我写了一个脚本，如下

import os


LABELS_PATH = "yolo_source_code/yolov10-main/runs/detect/predict2/labels"
RESULT_CSV_FILENAME = "使用train4的best模型预测结果.csv"
PEOPLE_LABEL_ID = [3]
VEHICLE_LABEL_ID = [0, 4, 5, 6, 7]

file_list = os.listdir(LABELS_PATH)

with open(RESULT_CSV_FILENAME, "w", encoding="utf-8") as f1:
    f1.write("image_name,people_num,vehicle_num\n")
    for each_label in file_list:
        people_num = 0
        vehicle_num = 0
        with open(f"{LABELS_PATH}/{each_label}", "r", encoding="utf-8") as f2:
            image_filename = each_label.replace(".txt", "")
            for line in f2:
                class_id = int(line.strip().split(" ")[0])
                if class_id in PEOPLE_LABEL_ID:
                    people_num += 1
                elif class_id in VEHICLE_LABEL_ID:
                    vehicle_num += 1
        f1.write(f"{image_filename},{people_num},{vehicle_num}\n")

只需设置LABELS_PATH和RESULT_CSV_FILENAME这两个变量值即可，第一个是预测后labels的路径，第二个是输出的csv文件名

执行脚本得到的csv文件格式如图

3. 检查模型在验证集、训练集上的准确性

可以用类似的方法检查模型在验证、训练集上的准确性，具体做法是用模型先预测一下训练集、验证集的图片，整理出csv文件

然后再把原本训练集、验证集的标注数据转化成一样的csv，然后看看完全相同的行有多少，看看准确率的百分比如何

五、写在最后

我不确定这个模型的最终效果如何，由于比赛并不透明公开成绩和答案，因此这个谁也说不准，重要的是在这个过程中学习知识吧

后续我准备想办法解决一下内存爆炸的问题，不行就加内存条看看行不行