MMDetection 训练实战指南:标准数据集与自定义数据集的完整训练流程
2026/9/20 1:29:45 网站建设 项目流程
  • 人工智能
  • 计算机视觉
  • 深度学习
  • 模型评测

【免费下载链接】mmdetection

OpenMMLab Detection Toolbox and Benchmark

项目地址:https://gitcode.com/gh_mirrors/mm/mmdetection
点击查看免费下载

MMDetection 为检测模型训练提供了开箱即用的完整工具链:以tools/train.py为统一入口,配合分布式脚本与灵活的配置体系,即可在 COCO 等标准数据集上训练预定义模型,也能通过数据集格式转换与配置文件定制,在自定义数据集上完成训练、测试与推理。读完本文,你将掌握学习率自动缩放、单卡/多卡/多机/Slurm 各类训练启动方式、resumeload-from的本质区别,以及将任意 VIA 格式标注转换为 COCO 格式并训练 Mask R-CNN 的完整实战方案。

一、训练工具链总览

MMDetection 的训练入口统一为 tools/train.py,其核心流程如下(源码可见 train.py):

  1. parse_args()解析命令行参数,包括配置文件路径、--work-dir--amp--auto-scale-lr--resume--cfg-options--launcher等;
  2. Config.fromfile(args.config)加载并解析配置文件;
  3. 若传入--cfg-options,通过cfg.merge_from_dict()将键值对合并进配置;
  4. 通过Runner.from_cfg(cfg)构建训练器,若配置中指定了runner_type,则从注册表RUNNERS.build(cfg)构建自定义 Runner;
  5. 调用runner.train()正式启动训练。

因此,一切训练行为都由配置文件驱动,命令行参数只做增量覆盖。理解这一点,是掌握后续所有训练方式的前提。

二、在标准数据集上训练预定义的模型

2.1 数据集准备

训练前需要先准备好数据集,具体的数据集下载、目录组织与标注说明请参考 数据集准备。

注意:目前configs/cityscapes文件夹下的配置文件均使用 COCO 预训练权重进行初始化。如果网络连接不可用或速度很慢,应提前下载现成的模型权重,否则训练刚开始时可能因为无法加载预训练权重而报错。

2.2 学习率自动缩放

背景:配置文件中的学习率是在 8 块 GPU、每块 GPU 2 张图像的设置下(批大小 = 8 × 2 = 16)确定的。该基准已在 configs/base/schedules/schedule_1x.py 中通过auto_scale_lr.base_batch_size = 16声明。学习率会以批大小 16 为基准自动缩放;同时,为了不影响其他基于 mmdet 的 codebase,auto_scale_lr.enable默认设置为False

启用方式:在启动命令中添加--auto-scale-lr参数即可。在启动前,请检查即将使用的配置文件名,因为配置名称通常指示默认批处理大小:

  • 默认批大小是 8 × 2 = 16,例如faster_rcnn_r50_caffe_fpn_90k_coco.pypisa_faster_rcnn_x101_32x4d_fpn_1x_coco.py
  • 非默认批次的配置文件名称中带有_NxM_字样,例如cornernet_hourglass104_mstest_32x3_210e_coco.py的批大小为 32 × 3 = 96,scnet_x101_64x4d_fpn_8x1_20e_coco.py的批大小为 8 × 1 = 8。

请记住:如果使用批大小不为 16 的配置文件,请检查配置底部是否有auto_scale_lr.base_batch_size;若找不到,可以在其继承的_base_ = [xxx]基配置文件中查找。想使用自动缩放学习率功能时,请不要修改这些值。

基本用法如下:

python tools/train.py \ ${CONFIG_FILE} \ --auto-scale-lr \ [optional arguments]

从源码看,--auto-scale-lr的底层逻辑位于 train.py:当该标志为真时,脚本检查配置中是否存在auto_scale_lrauto_scale_lr.enableauto_scale_lr.base_batch_size三个键,若存在则将cfg.auto_scale_lr.enable置为True,否则直接抛出RuntimeError提示配置缺失。

执行命令后,脚本会根据机器的 GPU 数量和训练批大小,按照线性扩展规则(Linear Scaling Rule,出自论文《Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour》)自动缩放学习率。例如:4 块 GPU、每块 GPU 2 张图片时lr=0.01,那么在 16 块 GPU、每块 GPU 4 张图片时,学习率会自动缩放至lr=0.08(批大小翻 8 倍,学习率随之放大 8 倍)。

如果不启用该功能,则需要按线性扩展规则手动计算并修改配置文件中optimizer.lr的值。以 configs/base/schedules/schedule_1x.py 为例,其默认优化器配置为SGD(lr=0.02, momentum=0.9, weight_decay=0.0001),若你的批大小翻倍,则需相应地将lr调为 0.04。

2.3 使用单 GPU 训练

tools/train.py支持在单张 GPU 上直接启动训练,基本用法如下:

python tools/train.py \ ${CONFIG_FILE} \ [optional arguments]

训练期间,日志文件和 checkpoint 文件会保存在工作目录下,该目录通过配置文件的work_dir字段或 CLI 参数--work-dir指定。从 train.py 源码可以确认工作目录的确定优先级为:CLI 参数 > 配置文件中的字段 > 配置文件名——若未显式指定,默认使用./work_dirs/${配置文件名}

默认情况下,模型会在每一轮训练后在验证集上评估一次,评估频率通过train_cfg配置:

# 每 12 轮迭代进行一次测试评估 train_cfg = dict(val_interval=12)

不同调度配置的默认值不同,例如 schedule_1x.py 中 1x 调度为max_epochs=12, val_interval=1,即每轮都验证。模型保存频率由default_hooks中的CheckpointHook控制,configs/base/default_runtime.py 默认interval=1(每轮保存一次 checkpoint)。

该工具接受以下参数:

  • --work-dir ${WORK_DIR}:覆盖工作目录。
  • --resume:自动从work_dir中的最新 checkpoint 恢复训练。
  • --resume ${CHECKPOINT_FILE}:从指定 checkpoint 文件继续训练。
  • --cfg-options 'Key=value':覆盖配置文件中的其他设置,例如--cfg-options 'dist_params.port=29500'。值若为列表需写成key="[a,b]"key=a,b,嵌套列表/元组可写成key="[(a,b),(c,d)]"(引号必须、不能有空格)。
  • --amp:启用自动混合精度训练。从 train.py 源码可见,该参数会将optim_wrapper.type改为AmpOptimWrapper并将loss_scale设为dynamic

注意resumeload-from的区别

  • resume既加载模型权重和优化器状态,也会继承指定 checkpoint 的迭代次数,训练不会重新开始。它作为命令行参数传入(--resume),源码中当--resumeauto时会设置cfg.resume = True且清空cfg.load_from;指定路径时则同时设置cfg.resume = Truecfg.load_from = 路径(见 train.py)。
  • load-from只加载模型权重,训练从头开始,常用于微调(fine-tune)。它需要写入配置文件中,而非命令行参数。default_runtime.py 中默认load_from = None, resume = False

2.4 使用 CPU 训练

CPU 训练的流程与单 GPU 完全一致,只需在训练前禁用 GPU:

export CUDA_VISIBLE_DEVICES=-1

之后直接运行单 GPU 训练脚本即可。

注意:官方不推荐使用 CPU 训练,因为速度过于缓慢;支持 CPU 仅是为了方便在没有 GPU 的机器上调试代码流程。

2.5 在多 GPU 上训练

tools/dist_train.sh用于启动多 GPU 训练,基本用法如下:

bash ./tools/dist_train.sh \ ${CONFIG_FILE} \ ${GPU_NUM} \ [optional arguments]

可选参数与单 GPU 训练一致。从 dist_train.sh 源码可以看出,其底层通过python -m torch.distributed.launch拉起分布式训练,NNODES(默认 1)、NODE_RANK(默认 0)、PORT(默认 29500)、MASTER_ADDR(默认 127.0.0.1)均作为环境变量读取,--nproc_per_node=$GPUS指定每节点进程数,--launcher pytorch被传给train.py

同时启动多个任务

如果想在一台机器上同时启动多个任务,例如在 8 块 GPU 的机器上启动 2 个各需 4 块 GPU 的任务,必须为不同任务指定不同的通信端口(默认 29500)以避免冲突。使用dist_train.sh时可通过环境变量设置端口:

CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 ./tools/dist_train.sh ${CONFIG_FILE} 4 CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 ./tools/dist_train.sh ${CONFIG_FILE} 4

2.6 使用多台机器训练

如果需要使用 ethernet 连接的多台机器协同训练,可以按如下方式在两台机器上分别启动:

在第一台机器上:

NNODES=2 NODE_RANK=0 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS

在第二台机器上:

NNODES=2 NODE_RANK=1 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS

其中MASTER_ADDR需指向第一台机器的 IP 地址,PORT为约定的通信端口。需要提醒的是:如果这几台机器之间没有使用高速网络连接,训练会非常慢,实践中应优先考虑 InfiniBand 等高速互联。

2.7 使用 Slurm 管理任务

Slurm 是常见的计算集群调度系统。在 Slurm 管理的集群上,可以使用 tools/slurm_train.sh 启动训练任务,它同时支持单节点与多节点训练。

基本用法如下:

[GPUS=${GPUS}] ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR}

以下是在名称为dev的 Slurm 分区上,使用 16 块 GPU 训练 Mask R-CNN,并将work-dir设置在共享文件系统下的示例:

GPUS=16 ./tools/slurm_train.sh dev mask_r50_1x configs/mask_rcnn_r50_fpn_1x_coco.py /nfs/xxxx/mask_rcnn_r50_fpn_1x

从 slurm_train.sh 源码可见,脚本支持的环境变量还包括GPUS_PER_NODE(默认 8)、CPUS_PER_TASK(默认 5)、SRUN_ARGS以及透传给训练脚本的PY_ARGS${@:5},即第 5 个参数起的全部参数),可通过srun--gres=gpu--ntasks等参数控制资源分配。

使用 Slurm 时,端口需通过以下两种方式之一设置:

方式一:通过--cfg-options设置端口(推荐)。这种方式无需改动原始配置文件:

CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR} --cfg-options 'dist_params.port=29500' CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR} --cfg-options 'dist_params.port=29501'

方式二:修改配置文件设置不同通信端口。在config1.py中设置:

dist_params = dict(backend='nccl', port=29500)

config2.py中设置:

dist_params = dict(backend='nccl', port=29501)

然后分别用两个配置启动任务:

CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR} CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR}

三、在自定义数据集上进行训练

本节以 balloon dataset(VIA 标注格式的气球数据集)为例,完整演示如何用自定义数据集对预定义模型进行训练、测试与推理。基本步骤如下:

  1. 准备自定义数据集;
  2. 准备配置文件;
  3. 在自定义数据集上进行训练、测试和推理。

3.1 支持新数据集的三种方式

MMDetection 一共支持三种形式应用新数据集:

  1. 将数据集重新组织为 COCO 格式;
  2. 将数据集重新组织为一个中间格式;
  3. 实现一个新的数据集类。

官方通常建议使用前两种方法,因为它们比实现新数据集类要简单得多。本文演示第一种:将数据转化为 COCO 格式。

注意:MMDetection 3.0 之后,数据集与评价指标已经解耦(除 CityScapes 外)。因此,用户在验证阶段可以用任意评价指标评价模型在任意数据集上的性能,例如用 VOC 评价指标评价模型在 COCO 数据集上的表现,或同时使用 VOC 与 COCO 两种评价指标评价 OpenImages 数据集上的模型。

3.2 COCO 标注格式

用于实例分割的 COCO 数据集格式如下,其中的键(key)都是必需的(更完整的字段细节可参考 COCO 官方数据格式说明):

{ "images": [image], "annotations": [annotation], "categories": [category] } image = { "id": int, "width": int, "height": int, "file_name": str, } annotation = { "id": int, "image_id": int, "category_id": int, "segmentation": RLE or [polygon], "area": float, "bbox": [x,y,width,height], # (x, y) 为 bbox 左上角的坐标 "iscrowd": 0 or 1, } categories = [{ "id": int, "name": str, "supercategory": str, }]

假设现在使用 balloon dataset。下载数据集后,需要实现一个函数将 VIA 标注格式转化为 COCO 格式,之后就能直接使用仓库中已实现的CocoDataset类(定义于 mmdet/datasets/coco.py)加载数据并训练与评测。

浏览 balloon 数据集会发现,其原始标注格式如下(每张图片的所有标注由 JSON 中所有键组成):

{'base64_img_data': '', 'file_attributes': {}, 'filename': '34020010494_e5cb88e1c4_k.jpg', 'fileref': '', 'regions': {'0': {'region_attributes': {}, 'shape_attributes': {'all_points_x': [1020, 1000, 994, 1003, 1023, 1050, 1089, 1134, 1190, 1265, 1321, 1361, 1403, 1428, 1442, 1445, 1441, 1427, 1400, 1361, 1316, 1269, 1228, 1198, 1207, 1210, 1190, 1177, 1172, 1174, 1170, 1153, 1127, 1104, 1061, 1032, 1020], 'all_points_y': [963, 899, 841, 787, 738, 700, 663, 638, 621, 619, 643, 672, 720, 765, 800, 860, 896, 942, 990, 1035, 1079, 1112, 1129, 1134, 1144, 1153, 1166, 1166, 1150, 1136, 1129, 1122, 1112, 1084, 1037, 989, 963], 'name': 'polygon'}}}, 'size': 1115004}

将 balloon dataset 转化为 COCO 格式的代码如下所示。核心思路是:逐张图片读取 VIA 标注,用mmcv.imread获取图片宽高,把all_points_x/all_points_y多边形点序列转换为 COCO 的segmentationbbox,最终通过mmengine.fileio.dump写出 COCO JSON 文件:

import os.path as osp import mmcv from mmengine.fileio import dump, load from mmengine.utils import track_iter_progress def convert_balloon_to_coco(ann_file, out_file, image_prefix): data_infos = load(ann_file) annotations = [] images = [] obj_count = 0 for idx, v in enumerate(track_iter_progress(data_infos.values())): filename = v['filename'] img_path = osp.join(image_prefix, filename) height, width = mmcv.imread(img_path).shape[:2] images.append( dict(id=idx, file_name=filename, height=height, width=width)) for _, obj in v['regions'].items(): assert not obj['region_attributes'] obj = obj['shape_attributes'] px = obj['all_points_x'] py = obj['all_points_y'] poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)] poly = [p for x in poly for p in x] x_min, y_min, x_max, y_max = (min(px), min(py), max(px), max(py)) data_anno = dict( image_id=idx, id=obj_count, category_id=0, bbox=[x_min, y_min, x_max - x_min, y_max - y_min], area=(x_max - x_min) * (y_max - y_min), segmentation=[poly], iscrowd=0) annotations.append(data_anno) obj_count += 1 coco_format_json = dict( images=images, annotations=annotations, categories=[{ 'id': 0, 'name': 'balloon' }]) dump(coco_format_json, out_file) if __name__ == '__main__': convert_balloon_to_coco(ann_file='data/balloon/train/via_region_data.json', out_file='data/balloon/train/annotation_coco.json', image_prefix='data/balloon/train') convert_balloon_to_coco(ann_file='data/balloon/val/via_region_data.json', out_file='data/balloon/val/annotation_coco.json', image_prefix='data/balloon/val')

使用如上函数,用户可以成功将 VIA 标注文件转化为 COCO JSON 格式,之后即可用CocoDataset加载数据训练模型,并用CocoMetric(定义于 mmdet/evaluation/metrics/coco_metric.py)进行评测。

3.3 准备配置文件

第二步是准备一个配置文件以成功加载数据集。假设我们要用 balloon dataset 训练一个配备 FPN 的 Mask R-CNN,将配置文件命名为mask-rcnn_r50-caffe_fpn_ms-poly-1x_balloon.py,保存在configs/balloon/目录下,内容如下。详细的配置编写方法可参考 学习配置文件。

# 新配置继承了基本配置,并做了必要的修改 _base_ = '../mask_rcnn/mask-rcnn_r50-caffe_fpn_ms-poly-1x_coco.py' # 我们还需要更改 head 中的 num_classes 以匹配数据集中的类别数 model = dict( roi_head=dict( bbox_head=dict(num_classes=1), mask_head=dict(num_classes=1))) # 修改数据集相关配置 data_root = 'data/balloon/' metainfo = { 'classes': ('balloon', ), 'palette': [ (220, 20, 60), ] } train_dataloader = dict( batch_size=1, dataset=dict( data_root=data_root, metainfo=metainfo, ann_file='train/annotation_coco.json', data_prefix=dict(img='train/'))) val_dataloader = dict( dataset=dict( data_root=data_root, metainfo=metainfo, ann_file='val/annotation_coco.json', data_prefix=dict(img='val/'))) test_dataloader = val_dataloader # 修改评价指标相关配置 val_evaluator = dict(ann_file=data_root + 'val/annotation_coco.json') test_evaluator = val_evaluator # 使用预训练的 Mask R-CNN 模型权重来做初始化,可以提高模型性能 load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'

对照仓库中的基配置 configs/mask_rcnn/mask-rcnn_r50-caffe_fpn_ms-poly-1x_coco.py 与 configs/base/datasets/coco_detection.py,可以逐项理解上述修改的含义:

  • _base_指向 COCO 上的 Mask R-CNN 训练配置(caffe 风格骨干 + FPN + 多尺度训练 + 多边形 mask),新配置只需在其上做增量修改;
  • roi_head中把bbox_headmask_headnum_classes从 80 改为 1,以匹配 balloon 数据集只有balloon一个类别的实际情况;
  • metainfo声明类名classes与可视化调色板palette,与CocoDatasetmetainfo机制对应;
  • train_dataloader/val_dataloader指定data_rootann_filedata_prefix(图片子目录)与batch_size(此处设为 1),test_dataloader直接复用val_dataloader
  • val_evaluator只需给定ann_file(类型沿用基配置中的CocoMetric),这正是 3.0 数据集与指标解耦的体现;
  • load_from指向 COCO 预训练权重,用于初始化模型参数以提升收敛速度与最终精度。

3.4 训练一个新的模型

使用上述新配置训练模型,只需运行:

python tools/train.py configs/balloon/mask-rcnn_r50-caffe_fpn_ms-poly-1x_balloon.py

更多详细用法(多卡、自动缩放学习率等)可回看本文"在标准数据集上训练预定义的模型"一节,二者完全通用。

3.5 测试以及推理

为测试训练完毕的模型,运行如下命令:

python tools/test.py configs/balloon/mask-rcnn_r50-caffe_fpn_ms-poly-1x_balloon.py work_dirs/mask-rcnn_r50-caffe_fpn_ms-poly-1x_balloon/epoch_12.pth

其中第二个参数为训练产出的 checkpoint 文件路径。测试与推理的更多详细用法(结果可视化、结果提交等)请参考 测试现有模型。

四、实战要点与常见坑位小结

  1. 先确认批大小再看学习率:配置文件命名中的_NxM_直接反映批大小;使用--auto-scale-lr前务必确认auto_scale_lr.base_batch_size存在(可能在继承的_base_中),且不要随意改动基准值,否则自动缩放会失去基准。
  2. resumeload-from别混淆:前者恢复训练(含优化器状态与迭代进度,命令行传入),后者仅加载权重重新开始(配置内写入,常用于微调)。
  3. 多任务必须错开端口:同机多任务或 Slurm 多任务共用默认端口 29500 会冲突,优先用--cfg-options 'dist_params.port=xxxx'临时指定,避免改动配置文件。
  4. 自定义数据集优先转 COCO 格式:这是最省力的接入方式,直接复用CocoDataset+CocoMetric;转换时务必保证segmentationbboxareaiscrowd等必需字段齐全。
  5. CPU 训练仅用于调试:通过CUDA_VISIBLE_DEVICES=-1禁用 GPU 后即可跑通全流程,但训练速度极慢,不应用于正式实验。

五、深入阅读

  • tools/train.py:训练入口,命令行参数解析与 Runner 构建
  • tools/dist_train.sh:多 GPU 分布式训练启动脚本
  • tools/slurm_train.sh:Slurm 集群训练启动脚本
  • configs/base/schedules/schedule_1x.py:1x 调度与auto_scale_lr基准定义
  • configs/base/datasets/coco_detection.py:COCO 数据集与评估器基配置
  • mmdet/datasets/coco.py:CocoDataset实现
  • mmdet/evaluation/metrics/coco_metric.py:CocoMetric实现
  • 数据集准备、学习配置文件、测试现有模型
  • 人工智能
  • 计算机视觉
  • 深度学习
  • 模型评测

【免费下载链接】mmdetection

OpenMMLab Detection Toolbox and Benchmark

项目地址:https://gitcode.com/gh_mirrors/mm/mmdetection
点击查看免费下载

相关推荐

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询