1. AI技术聚合首页
  2. Python

MMYOLO 自定义数据集从标注到部署保姆级教程

青葱年少 • • Python • 阅读 100

theme: juejin

来自社区 PeterH0323 投稿

AI 已经被应用到各行各业，现如今任何人都可以轻松基于开源框架快速搭建符合自身需求的 AI 应用。本文将基于 MMYOLO 开源框架，基于生活中收集的猫猫数据集，教你如何从零开始训练一个可部署检测模型。

本文档配套的视频已发布在 B 站，可前去查看: 自定义数据集从标注到部署保姆级教程

本教程所有指令是在 Linux 上面完成，Windows 也是完全可用的，但是命令和操作稍有不同。

本文为实践操作类教程。为了让不同水平用户都能基于自定义数据训练出一个不错的模型，本文将分成多个步骤详细描述。如果你自定义数据集训练效果不佳，麻烦你按照本文步骤逐条确认下是否有不正确或者不合理的地方。

完整步骤如下

  1. 数据集准备：tools/misc/download_dataset.py
  2. 使用 labelme 和算法进行辅助标注：demo/image_demo.py + labelme
  3. 使用脚本转换成 COCO 数据集格式：tools/dataset_converters/labelme2coco.py
  4. 数据集划分为训练集、验证集和测试集：tools/misc/coco_split.py
  5. 构建 config 文件
  6. 数据集可视化分析：tools/analysis_tools/dataset_analysis.py
  7. 优化 anchor 尺寸：tools/analysis_tools/optimize_anchors.py
  8. 可视化数据处理部分： tools/analysis_tools/browse_dataset.py
  9. 启动训练：tools/train.py
  10. 模型推理：demo/image_demo.py
  11. 模型部署

在训练得到模型权重和验证集的 mAP 后，用户需要对预测错误的 bad case 进行深入分析，以便优化模型，MMYOLO 在后续会增加这个功能，敬请期待。

由于内容较多，因此本文只描述核心部分，详细文档请看视频或者官方文档： https://github.com/open-mmlab/mmyolo/blob/main/docs/zh_cn/user_guides/custom_dataset.md

本教程基于 MMYOLO 0.2.0，随着版本迭代更新，可能有些参数会有所改变，如果发现就不一致则请切换到 dev 分支使用，请知悉！

1. 数据集准备

为了方便大家从零开始并且不用自己准备数据集，MMYOLO 热心社区用户 @RangeKing 提供了一个 144 张图片的 cat 数据集，并由 @PeterH0323 进行数据清洗，本文都以这个小数据集为例进行讲解。

下载也非常简单，只需要一条命令即可完成（数据集压缩包大小 217 MB）：

python tools/misc/download_dataset.py --dataset-name cat --save-dir ./data/cat --unzip --delete

该命令会自动下载数据集到 ./data/cat 文件夹中，该文件的目录结构是：

.
└── ./data/cat
    ├── images # 图片文件
    │    ├── image1.jpg
    │    ├── image2.png
    │    └── ...
    ├── labels # labelme 标注文件
    │    ├── image1.json
    │    ├── image2.json
    │    └── ...
    ├── annotations # 数据集划分的 COCO 文件
    │    ├── annotations_all.json # 全量数据的 COCO label 文件
    │    ├── trainval.json # 划分比例 80% 的数据
    │    └── test.json # 划分比例 20% 的数据
    └── class_with_id.txt # id + class_name 文件

这个数据集可以直接训练，如果您想体验整个流程的话，可以将 images 文件夹以外的其余文件都删除。

2. 使用 labelme 和算法进行辅助标注

有了数据集后，如果要训练必然是要标注的，通常，标注有 2 种方法：

  • 软件或者算法辅助 + 人工修正 label（推荐，降本提速）
  • 仅人工标注

为了减少标注时间，MMYOLO 特意基于 labelme 支持了软件辅助标注功能。

辅助标注的原理是用已有模型进行推理，将得出的推理信息保存为标注软件 label 文件格式。然后人工操作标注软件加载生成好的 label 文件，只需要检查每张图片的目标是否标准，以及是否有漏掉、错标的目标。

如果已有模型（典型的如 COCO 预训练模型）没有您自定义新数据集的类别，建议先人工打 100 张左右的图片 label，训练个初始模型，然后再进行辅助标注。

下面会分别介绍其过程。

2.1 算法辅助标注

使用 MMYOLO 提供的模型推理脚本 demo/image_demo.py，并设置 --to-labelme 则可以将推理结果生成 labelme 格式的 label 文件。

这里使用 YOLOv5-s 作为例子来进行辅助标注刚刚下载的 cat 数据集，先下载 YOLOv5-s 的权重:

mkdir work_dirs
wget https://download.openmmlab.com/mmyolo/v0/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth -P ./work_dirs

由于 COCO 80 类数据集中已经包括了 cat 这一类，因此我们可以直接加载 COCO 预训练权重进行辅助标注：

python demo/image_demo.py ./data/cat/images \
                          ./configs/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py \
                          ./work_dirs/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth \
                          --out-dir ./data/cat/labels \
                          --class-name cat \
                          --to-labelme

生成的 label 文件会在 --out-dir 中:

.
└── $OUT_DIR
    ├── image1.json
    ├── image1.json
    └── ...

这是一张原图及其生成的 json 例子：

2.2 人工校正标签

在算法辅助标注后，需要人工检查伪标签是否有问题，本教程使用的标注软件是 labelme。

  • 安装 labelme
conda create -n labelme python=3.8
conda activate labelme
pip install labelme==5.1.1
  • 启动 labelme
cd /path/to/mmyolo
labelme ./data/cat/images --output ./data/cat/labels --autosave --nodata

输入命令之后 labelme 就会启动，然后进行 label 检查即可。如果 labelme 启动失败，命令行输入 export QT_DEBUG_PLUGINS=1 查看具体缺少什么库，安装一下即可。下面是一张采用 labelme 软件打开的效果图：

MMYOLO 自定义数据集从标注到部署保姆级教程

3. 使用脚本转换成 COCO 数据集格式

标注的格式需要转换为最常用的 COCO 格式，这也是方便后续可视化分析。

3.1 使用脚本转换

MMYOLO 提供脚本将 labelme 的 label 转换为 COCO label：

python tools/dataset_converters/labelme2coco.py --img-dir ${图片文件夹路径} \
                                                --labels-dir ${label 文件夹位置} \
                                                --out ${输出 COCO label json 路径} \
                                                [--class-id-txt ${class_with_id.txt 路径}]

其中： --class-id-txt：是数据集 id class_name 的 .txt 文件：

  • 如果不指定，则脚本会自动生成，生成在 --out 同级的目录中，保存文件名为 class_with_id.txt；
  • 如果指定，脚本仅会进行读取但不会新增或者覆盖，同时，脚本里面还会判断是否存在 .txt 中其他的类，如果出现了会报错提示，届时，请用户检查 .txt 文件并加入新的类及其 id。

.txt 文件的例子如下（ id 可以和 COCO 一样，从 1 开始）：

1 cat
2 dog
3 bicycle
4 motorcycle

以本教程的 cat 数据集为例：

python tools/dataset_converters/labelme2coco.py --img-dir ./data/cat/images \
                                                --labels-dir ./data/cat/labels \
                                                --out ./data/cat/annotations/annotations_all.json

本次演示的 cat 数据集（注意不需要包括背景类），可以看到生成的 class_with_id.txt 中只有 1 类：

1 cat

3.2 检查转换的 COCO label

使用下面的命令可以将 COCO 的 label 在图片上进行显示，这一步可以验证刚刚转换是否有问题：

python tools/analysis_tools/browse_coco_json.py --img-dir ./data/cat/images \
                                                --ann-file ./data/cat/annotations/annotations_all.json

关于 tools/analysis_tools/browse_coco_json.py 的更多用法请参考 可视化 COCO label。

4. 数据集划分为训练集、验证集和测试集

通常，自定义图片都是一个大文件夹，里面全部都是图片，需要我们自己去对图片进行训练集、验证集、测试集的划分，如果数据量比较少，可以不划分验证集。下面是划分脚本的具体用法：

python tools/misc/coco_split.py --json ${COCO label json 路径} \
                                --out-dir ${划分 label json 保存根路径} \
                                --ratios ${划分比例} \
                                [--shuffle] \
                                [--seed ${划分的随机种子}]

其中：

  • --ratios：划分的比例，如果只设置了 2 个，则划分为 trainval + test；如果设置为 3 个，则划分为 train + val + test。支持两种格式 —— 整数、小数：

    • 整数：按比例进行划分，代码中会进行归一化之后划分数据集。例子： --ratio 2 1 1（代码里面会转换成 0.5 0.25 0.25） or --ratio 3 1（代码里面会转换成 0.75 0.25）
    • 小数：划分为比例。如果加起来不为 1 ，则脚本会进行自动归一化修正。例子： --ratio 0.8 0.1 0.1 or --ratio 0.8 0.2

  • --shuffle: 是否打乱数据集再进行划分。

  • --seed：设定划分的随机种子，不设置的话自动生成随机种子。

例子：

python tools/misc/coco_split.py --json ./data/cat/annotations/annotations_all.json \
                                --out-dir ./data/cat/annotations \
                                --ratios 0.8 0.2 \
                                --shuffle \
                                --seed 10

5. 构建 config 文件

在自定义的数据集准备工作都完成后，需要自己新建一个 config 并加入需要修改的部分信息。此处以 YOLOv5-s 算法为例，config 的命名规则：

  • config 继承的是 yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py；
  • 训练的类以本教程提供的数据集中的类 cat 为例（如果是自己的数据集，可以自定义类型的总称）；
  • 本教程测试的显卡型号是 1 x 3080Ti 12G 显存，电脑内存 32G，可以训练 YOLOv5-s 最大批次是 batch size = 32（详细机器资料可见附录）；
  • 训练轮次是 100 epoch。

综上所述：可以将其命名为 yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py，并将其放置在文件夹 configs/custom_dataset 中。

我们可以在 configs 目录下新建一个新的目录 custom_dataset，同时在里面新建该 config 文件，并添加以下内容：

_base_ = '../yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py'

max_epochs = 100  # 训练的最大 epoch
data_root = './data/cat/'  # 数据集目录的绝对路径
# data_root = '/root/workspace/mmyolo/data/cat/'  # Docker 容器里面数据集目录的绝对路径

# 结果保存的路径，可以省略，省略保存的文件名位于 work_dirs 下 config 同名的文件夹中
# 如果某个 config 只是修改了部分参数，修改这个变量就可以将新的训练文件保存到其他地方
work_dir = './work_dirs/yolov5_s-v61_syncbn_fast_1xb32-100e_cat'

# load_from 可以指定本地路径或者 URL，设置了 URL 会自动进行下载，因为上面已经下载过，我们这里设置本地路径
# 因为本教程是在 cat 数据集上微调，故这里需要使用 `load_from` 来加载 MMYOLO 中的预训练模型，这样可以在加快收敛速度的同时保证精度
load_from = './work_dirs/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth'  # noqa

# 根据自己的 GPU 情况，修改 batch size，YOLOv5-s 默认为 8卡 x 16bs
train_batch_size_per_gpu = 32
train_num_workers = 4  # 推荐使用 train_num_workers = nGPU x 4

save_epoch_intervals = 2  # 每 interval 轮迭代进行一次保存一次权重

# 根据自己的 GPU 情况，修改 base_lr，修改的比例是 base_lr_default * (your_bs / default_bs)
base_lr = _base_.base_lr / 4

anchors = [  # 此处已经根据数据集特点更新了 anchor，关于 anchor 的生成，后面小节会讲解
    [(68, 69), (154, 91), (143, 162)],  # P3/8
    [(242, 160), (189, 287), (391, 207)],  # P4/16
    [(353, 337), (539, 341), (443, 432)]  # P5/32
]

class_name = ('cat', )  # 根据 class_with_id.txt 类别信息，设置 class_name
num_classes = len(class_name)
metainfo = dict(
    CLASSES=class_name, # 注意：这个字段在最新版本中换成了小写
    PALETTE=[(220, 20, 60)]  # 画图时候的颜色，随便设置即可
)

train_cfg = dict(
    max_epochs=max_epochs,
    val_begin=20,  # 第几个 epoch 后验证，这里设置 20 是因为前 20 个 epoch 精度不高，测试意义不大，故跳过
    val_interval=save_epoch_intervals  # 每 val_interval 轮迭代进行一次测试评估
)

model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        prior_generator=dict(base_sizes=anchors),

        # loss_cls 会根据 num_classes 动态调整，但是 num_classes = 1 的时候，loss_cls 恒为 0
        loss_cls=dict(loss_weight=0.5 *
                      (num_classes / 80 * 3 / _base_.num_det_layers))))

train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        _delete_=True,
        type='RepeatDataset',
        # 数据量太少的话，可以使用 RepeatDataset ，在每个 epoch 内重复当前数据集 n 次，这里设置 5 是重复 5 次
        times=5,
        dataset=dict(
            type=_base_.dataset_type,
            data_root=data_root,
            metainfo=metainfo,
            ann_file='annotations/trainval.json',
            data_prefix=dict(img='images/'),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=_base_.train_pipeline)))

val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='annotations/trainval.json',
        data_prefix=dict(img='images/')))

test_dataloader = val_dataloader

val_evaluator = dict(ann_file=data_root + 'annotations/trainval.json')
test_evaluator = val_evaluator

optim_wrapper = dict(optimizer=dict(lr=base_lr))

default_hooks = dict(
    # 设置间隔多少个 epoch 保存模型，以及保存模型最多几个，`save_best` 是另外保存最佳模型（推荐）
    checkpoint=dict(
        type='CheckpointHook',
        interval=save_epoch_intervals,
        max_keep_ckpts=5,
        save_best='auto'),
    param_scheduler=dict(max_epochs=max_epochs),
    # logger 输出的间隔
    logger=dict(type='LoggerHook', interval=10))

我们在 projects/misc/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 放了一份相同的 config 文件，用户可以选择复制到 configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 路径直接开始训练。

6. 数据集可视化分析

数据集可视化分析有两个目的：

  1. 验证配置文件写的是否正确
  2. 对数据集有个初步分析

脚本 tools/analysis_tools/dataset_analysis.py 能够帮助用户得到数据集的分析图。该脚本可以生成 4 种分析图：

  • 显示类别和 bbox 实例个数的分布图：show_bbox_num
  • 显示类别和 bbox 实例宽、高的分布图：show_bbox_wh
  • 显示类别和 bbox 实例宽/高比例的分布图：show_bbox_wh_ratio
  • 基于面积规则下，显示类别和 bbox 实例面积的分布图：show_bbox_area

以本教程 cat 数据集 的 config 为例。

查看训练集数据分布情况：

python tools/analysis_tools/dataset_analysis.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \
                                                --out-dir work_dirs/dataset_analysis_cat/train_dataset

查看验证集数据分布情况：

python tools/analysis_tools/dataset_analysis.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \
                                                --out-dir work_dirs/dataset_analysis_cat/val_dataset \
                                                --val-dataset
基于面积规则下，显示类别和 bbox 实例面积的分布图显示类别和 bbox 实例宽、高的分布图
MMYOLO 自定义数据集从标注到部署保姆级教程MMYOLO 自定义数据集从标注到部署保姆级教程
显示类别和 bbox 实例个数的分布图显示类别和 bbox 实例宽/高比例的分布图
MMYOLO 自定义数据集从标注到部署保姆级教程MMYOLO 自定义数据集从标注到部署保姆级教程

基于面积规则下，显示类别和 bbox 实例面积的分布图

显示类别和 bbox 实例宽、高的分布图

显示类别和 bbox 实例个数的分布图

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-tmokTAs8-1672831496747)(https://p3-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/2e539f094b66445d91013458838e97fc~tplv-k3u1fbpfcp-zoom-1.image)]

显示类别和 bbox 实例宽/高比例的分布图

因为本教程使用的 cat 数据集数量比较少，故 config 里面用了 RepeatDataset，显示的数目实际上都是重复了 5 次。如果您想得到无重复的分析结果，可以暂时将 RepeatDataset 下面的 times 参数从 5 改成 1。

经过输出的图片分析可以得出，本教程使用的 cat 数据集的训练集具有以下情况：

  • 图片全部是 large object；
  • 类别 cat 的数量是 655；
  • bbox 的宽高比例大部分集中在 1.0 ~ 1.11，比例最小值是 0.36，最大值是 2.9；
  • bbox 的宽大部分是 500 ~ 600 左右，高大部分是 500 ~ 600 左右。

关于 tools/analysis_tools/dataset_analysis.py 的更多用法请参考 可视化数据集分析 。

7. 优化 Anchor 尺寸

YOLOv5 是一个 anchor-based 算法，需要对 anchor 进行自适应计算，anchor-free 的模型可以跳过此步骤。

脚本 tools/analysis_tools/optimize_anchors.py 支持 YOLO 系列中三种锚框生成方式，分别是 k-means、Differential Evolution、v5-k-means。

本示例使用的是 YOLOv5 进行训练，使用的是 640 x 640 的输入大小，使用 v5-k-means 进行描框的优化：

python tools/analysis_tools/optimize_anchors.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \
                                                --algorithm v5-k-means \
                                                --input-shape 640 640 \
                                                --prior-match-thr 4.0 \
                                                --out-dir work_dirs/dataset_analysis_cat

经过计算的 Anchor 如下：

修改 config 文件里面的 anchors 变量：

anchors = [
    [(68, 69), (154, 91), (143, 162)],  # P3/8
    [(242, 160), (189, 287), (391, 207)],  # P4/16
    [(353, 337), (539, 341), (443, 432)]  # P5/32
]

因为该命令使用的是 k-means 聚类算法，存在一定的随机性，这与初始化有关。故每次执行得到的 Anchor 都会有些不一样，但是都是基于传递进去的数据集来进行生成的，故不会有什么不良影响。

关于 tools/analysis_tools/optimize_anchors.py 的更多用法请参考 优化锚框尺寸。

8. 可视化数据处理部分

脚本 tools/analysis_tools/browse_dataset.py 能够帮助用户直接可视化 config 配置中数据处理部分，同时可以选择保存可视化图片到指定文件夹内。

下面演示使用我们刚刚新建的 config 文件 configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 来可视化图片，该命令会使得图片直接弹出显示，每张图片持续 3 秒，图片不进行保存：

python tools/analysis_tools/browse_dataset.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \
                                              --show-interval 3

关于 tools/analysis_tools/browse_dataset.py 的更多用法请参考 可视化数据集 。

9. 启动训练

下面会从以下 3 点来进行讲解：

  1. 训练可视化
  2. YOLOv5 模型训练
  3. 切换 YOLO 模型训练

9.1 训练可视化

如果需要采用浏览器对训练过程可视化，MMYOLO 目前提供 2 种方式： WandB 和 TensorBoard，根据自己的情况选择其一即可（后续会扩展更多可视化后端支持）。

9.1.1 WandB

WandB 可视化需要在官网注册，并在 https://wandb.ai/settings 获取到 WandB 的 API Keys。

然后在命令行进行安装：

pip install wandb
# 运行了 wandb login 后输入上文中获取到的 API Keys ，便登录成功。
wandb login

在我们刚刚新建的 config 文件 configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 的最后添加 WandB 配置：

visualizer = dict(vis_backends=[dict(type='LocalVisBackend'), dict(type='WandbVisBackend')])

9.1.2 TensorBoard

安装 Tensorboard 环境：

pip install tensorboard

在我们刚刚新建的 config 文件 configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 中添加 tensorboard 配置：

visualizer = dict(vis_backends=[dict(type='LocalVisBackend'),dict(type='TensorboardVisBackend')])

运行训练命令后，Tensorboard 文件会生成在可视化文件夹 work_dirs/yolov5_s-v61_syncbn_fast_1xb32-100e_cat/${TIMESTAMP}/vis_data 下， 运行下面的命令便可以在网页链接使用 Tensorboard 查看 loss、学习率和 coco/bbox_mAP 等可视化数据了：

tensorboard --logdir=work_dirs/yolov5_s-v61_syncbn_fast_1xb32-100e_cat

9.2 执行训练

使用下面命令进行启动训练（训练大约需要 2.5 个小时）：

python tools/train.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py

如果您开启了 WandB 的话，可以登录到自己的账户，在 WandB 中查看本次训练的详细信息了：

下面是 1 x 3080Ti、batch size = 32，训练 100 epoch 最佳精度权重 work_dirs/yolov5_s-v61_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_98.pth 得出来的精度（详细机器资料可见附录）：

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.968
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 1.000
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.968
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.886
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.977
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.977
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.977

bbox_mAP_copypaste: 0.968 1.000 1.000 -1.000 -1.000 0.968
Epoch(val) [98][116/116]  coco/bbox_mAP: 0.9680  coco/bbox_mAP_50: 1.0000  coco/bbox_mAP_75: 1.0000  coco/bbox_mAP_s: -1.0000  coco/bbox_mAP_m: -1.0000  coco/bbox_mAP_l: 0.9680

在一般的 finetune 最佳实践中都会推荐将 backbone 固定不参与训练，并且学习率 lr 也进行相应缩放，但是在本教程中发现这种做法会出现一定程度掉点。猜测可能原因是 cat 类别已经在 COCO 数据集中，而本教程使用的 cat 数据集数量比较小导致的。

下表是采用 MMYOLO YOLOv5 预训练模型 yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth 在没对 cat 数据集进行 finetune 的测试精度，可以看到 cat 类别的 mAP 只有 0.866，经过我们 finetune mAP 提升到了 0.968，提升了 10.2 %，可以证明训练是非常成功的：

+---------------+-------+--------------+-----+----------------+------+
| category      | AP    | category     | AP  | category       | AP   |
+---------------+-------+--------------+-----+----------------+------+
| person        | nan   | bicycle      | nan | car            | nan  |
| motorcycle    | nan   | airplane     | nan | bus            | nan  |
| train         | nan   | truck        | nan | boat           | nan  |
| traffic light | nan   | fire hydrant | nan | stop sign      | nan  |
| parking meter | nan   | bench        | nan | bird           | nan  |
| cat           | 0.866 | dog          | nan | horse          | nan  |
| sheep         | nan   | cow          | nan | elephant       | nan  |
| bear          | nan   | zebra        | nan | giraffe        | nan  |
| backpack      | nan   | umbrella     | nan | handbag        | nan  |
| tie           | nan   | suitcase     | nan | frisbee        | nan  |
| skis          | nan   | snowboard    | nan | sports ball    | nan  |
| kite          | nan   | baseball bat | nan | baseball glove | nan  |
| skateboard    | nan   | surfboard    | nan | tennis racket  | nan  |
| bottle        | nan   | wine glass   | nan | cup            | nan  |
| fork          | nan   | knife        | nan | spoon          | nan  |
| bowl          | nan   | banana       | nan | apple          | nan  |
| sandwich      | nan   | orange       | nan | broccoli       | nan  |
| carrot        | nan   | hot dog      | nan | pizza          | nan  |
| donut         | nan   | cake         | nan | chair          | nan  |
| couch         | nan   | potted plant | nan | bed            | nan  |
| dining table  | nan   | toilet       | nan | tv             | nan  |
| laptop        | nan   | mouse        | nan | remote         | nan  |
| keyboard      | nan   | cell phone   | nan | microwave      | nan  |
| oven          | nan   | toaster      | nan | sink           | nan  |
| refrigerator  | nan   | book         | nan | clock          | nan  |
| vase          | nan   | scissors     | nan | teddy bear     | nan  |
| hair drier    | nan   | toothbrush   | nan | None           | None |
+---------------+-------+--------------+-----+----------------+------+

9.3 尝试 MMYOLO 其他模型

MMYOLO 集成了多种 YOLO 算法，切换非常方便，无需重新熟悉一个新的 repo，直接切换 config 文件就可以轻松切换 YOLO 模型，只需简单 3 步即可切换模型：

  1. 构建 config 文件
  2. 下载预训练权重
  3. 启动训练

MMEngine 的配置文件写法非常灵活，支持多层继承，选择不同的 base 文件后继承重写的配置也不一样。以 YOLOv6-s 为例，如果你选择直接继承上述的 YOLOv5 cat 配置，那么需要改动的地方非常少。不过为了减少继承层级，我们依然以 YOLOv6-s 配置作为 base。

  1. 搭建一个新的 config：
_base_ = '../yolov6/yolov6_s_syncbn_fast_8xb32-400e_coco.py'

max_epochs = 100  # 训练的最大 epoch
data_root = './data/cat/'  # 数据集目录的绝对路径

# 结果保存的路径，可以省略，省略保存的文件名位于 work_dirs 下 config 同名的文件夹中
# 如果某个 config 只是修改了部分参数，修改这个变量就可以将新的训练文件保存到其他地方
work_dir = './work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat'

# load_from 可以指定本地路径或者 URL，设置了 URL 会自动进行下载，因为上面已经下载过，我们这里设置本地路径
# 因为本教程是在 cat 数据集上微调，故这里需要使用 `load_from` 来加载 MMYOLO 中的预训练模型，这样可以在加快收敛速度的同时保证精度
load_from = './work_dirs/yolov6_s_syncbn_fast_8xb32-400e_coco_20221102_203035-932e1d91.pth'  # noqa

# 根据自己的 GPU 情况，修改 batch size，YOLOv6-s 默认为 8卡 x 32bs
train_batch_size_per_gpu = 32
train_num_workers = 4  # 推荐使用 train_num_workers = nGPU x 4

save_epoch_intervals = 2  # 每 interval 轮迭代进行一次保存一次权重

# 根据自己的 GPU 情况，修改 base_lr，修改的比例是 base_lr_default * (your_bs / default_bs)
base_lr = _base_.base_lr / 8

class_name = ('cat', )  # 根据 class_with_id.txt 类别信息，设置 class_name
num_classes = len(class_name)
metainfo = dict(
    CLASSES=class_name, # 注意： 在新版本中已经换成了小写
    PALETTE=[(220, 20, 60)]  # 画图时候的颜色，随便设置即可
)

train_cfg = dict(
    max_epochs=max_epochs,
    val_begin=20,  # 第几个 epoch 后验证，这里设置 20 是因为前 20 个 epoch 精度不高，测试意义不大，故跳过
    val_interval=save_epoch_intervals,  # 每 val_interval 轮迭代进行一次测试评估
    dynamic_intervals=[(max_epochs - _base_.num_last_epochs, 1)]
)

model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes)),
    train_cfg=dict(
        initial_assigner=dict(num_classes=num_classes),
        assigner=dict(num_classes=num_classes))
)

train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        _delete_=True,
        type='RepeatDataset',
        # 数据量太少的话，可以使用 RepeatDataset ，在每个 epoch 内重复当前数据集 n 次，这里设置 5 是重复 5 次
        times=5,
        dataset=dict(
            type=_base_.dataset_type,
            data_root=data_root,
            metainfo=metainfo,
            ann_file='annotations/trainval.json',
            data_prefix=dict(img='images/'),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=_base_.train_pipeline)))

val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='annotations/trainval.json',
        data_prefix=dict(img='images/')))

test_dataloader = val_dataloader

val_evaluator = dict(ann_file=data_root + 'annotations/trainval.json')
test_evaluator = val_evaluator

optim_wrapper = dict(optimizer=dict(lr=base_lr))

default_hooks = dict(
    # 设置间隔多少个 epoch 保存模型，以及保存模型最多几个，`save_best` 是另外保存最佳模型（推荐）
    checkpoint=dict(
        type='CheckpointHook',
        interval=save_epoch_intervals,
        max_keep_ckpts=5,
        save_best='auto'),
    param_scheduler=dict(max_epochs=max_epochs),
    # logger 输出的间隔
    logger=dict(type='LoggerHook', interval=10))

custom_hooks = [
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0001,
        update_buffers=True,
        strict_load=False,
        priority=49),
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=max_epochs - _base_.num_last_epochs,
        switch_pipeline=_base_.train_pipeline_stage2)
]

同样，我们在 projects/misc/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py 放了一份相同的 config 文件，用户可以选择复制到 configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py 路径直接开始训练。 虽然新的 config 看上去好像很多东西，其实很多都是重复的，用户可以用对比软件对比一下即可看出大部分的配置都是和 yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 相同的。因为这 2 个 config 文件需要继承不同的 config，所以还是要添加一些必要的配置。

  1. 下载 YOLOv6-s 的预训练权重
wget https://download.openmmlab.com/mmyolo/v0/yolov6/yolov6_s_syncbn_fast_8xb32-400e_coco/yolov6_s_syncbn_fast_8xb32-400e_coco_20221102_203035-932e1d91.pth -P work_dirs/
  1. 训练
python tools/train.py configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py

在我的实验中，最佳模型是 work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth，其精度如下：

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.987
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 1.000
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.987
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.895
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.989
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.989
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.989

bbox_mAP_copypaste: 0.987 1.000 1.000 -1.000 -1.000 0.987
Epoch(val) [96][116/116]  coco/bbox_mAP: 0.9870  coco/bbox_mAP_50: 1.0000  coco/bbox_mAP_75: 1.0000  coco/bbox_mAP_s: -1.0000  coco/bbox_mAP_m: -1.0000  coco/bbox_mAP_l: 0.9870

以上演示的是如何在 MMYOLO 中切换模型，可以快速对不同模型进行精度对比，精度高的模型可以上线生产。在我的实验中，YOLOv6 最佳精度 0.9870 比 YOLOv5 最佳精度 0.9680 高出 1.9 %，故后续我们使用 YOLOv6 来进行讲解。

10. 模型推理

使用最佳的模型进行推理，下面命令中的最佳模型路径是 ./work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth，请用户自行修改为自己训练的最佳模型路径。

python demo/image_demo.py ./data/cat/images \
                          ./configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py \
                          ./work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth \
                          --out-dir ./data/cat/pred_images

如果推理结果不理想，这里举例 2 种情况： 1. 模型欠拟合： 需要先判断是不是训练 epoch 不够导致的欠拟合，如果是训练不够，则修改 config 文件里面的 max_epochs 和 work_dir 参数，或者根据上面的命名方式新建一个 config 文件，重新进行训练。 2. 数据集需优化： 如果 epoch 加上去了还是不行，可以增加数据集数量，同时可以重新检查并优化数据集的标注，然后重新进行训练。

11. 模型部署

MMYOLO 提供两种部署方式：

  1. MMDeploy 框架进行部署
  2. 使用 projects/easydeploy 进行部署

考虑到本部分细节较多，故不在本文说明，请跳转视频或者官方文档查看。

总结

本文从零开始教你如何基于 MMYOLO 训练一个可部署的检测模型，提供了 11 个完整步骤。用户在面对自定义数据集时候也可以按照这个步骤来逐条确认检查。

如果你对本文有任何好的建议，或者对 MMYOLO 有不同看法，欢迎留言反馈，或者直接联系微信小助手。

MMYOLO 官方地址： https://github.com/open-mmlab/mmyolo

文章出处登录后可见！

已经登录？立即刷新

共计人评分，平均分

到目前为止还没有投票！成为第一位评论此文章。

赞 (0)
青葱年少的头像青葱年少普通用户
0
上一篇 2023年9月13日
下一篇 2023年9月13日

相关推荐