mmdetection自定义数据集目标检测实战:从标注到训练避坑指南
最近在做一个工业场景的缺陷检测项目需要识别几种很偏门的瑕疵类型这类需求没有现成的预训练模型可以直接抄答案只能自己标注一批图片再用 mmdetection 训练一个自定义数据集的目标检测模型。前后折腾了两周踩了大大小小一堆坑最后总算跑通了。这篇文章就把整个流程重新梳理一遍从数据集制作、格式转换到训练评估、报错排查尽量覆盖你在实战中最可能遇到的每一个环节。mmdetection 是目前目标检测领域覆盖面最广的开源工具库之一Faster R-CNN、RetinaNet、YOLOv3、SSD、Cascade R-CNN 这类经典算法都能开箱即用而且配置文件设计得很统一替换数据集和网络结构基本就是改改字典的事。也正因为它管的算法多版本依赖特别复杂新手最容易被卡住的就是环境问题。这篇博文适合三类人第一次接触 mmdetection 想做自定义数据集的新手数据已经标好但不知道怎么组织成 COCO 格式的开发者以及训练中途老是报错、连日志都看不懂的同学。1. 项目整体思路拆解动手之前先想明白这四件事1.1 为什么选 mmdetection 而不是自己搭网络或换其他框架我在选框架时其实犹豫过一阵。如果用现成的 YOLOv5确实上手最快数据格式也简单但当时我需要对比 Faster R-CNN 和 Cascade R-CNN 多种算法还想后续试试不同 backbone 的效果。mmdetection 一个库就全包了换成不同算法直接换配置文件就行这点比每个模型单独搭一套工程舒服太多。这也是我最终选它的核心原因不是因为它最简单而是因为它最能满足“算法灵活对比”的需求。另一个原因是社区资料丰富。毕竟这个库用的人多真遇到报错搜一下基本能找到解决方案。对自定义数据集的训练流程官方文档也给出了很标准的路子把数据转成 COCO 格式然后改配置里的类别数和路径就能启动训练。这个思路看起来简单但细节藏得深后面我会把每一步拆开讲尤其是那些不改就会突然给你罢工的坑。1.2 自定义数据集的核心流程梳理整个项目可以拆成四条主线环境搭建、数据准备、训练评估、问题排查。环境搭建解决的是“代码能不能跑起来”数据准备解决的是“模型能看到什么”训练评估解决的是“模型学得好不好”问题排查则是贯穿全程的兜底能力。别小看“模型能看到什么”这半句话。目标检测模型的训练质量不只是网络结构决定的数据标注是否规范、类别定义是否统一、训练集和验证集划分是否合理都会直接影响最终 mAP。我第一次训练时标注不严谨一个类别名称在部分图片里叫 defect另一些叫 Defect结果模型学得稀里糊涂loss 降不下去最后查了半天才发现是这种低级问题。所以从头开始就规范起来反而最省时间。1.3 先规划再动手数据量、目录结构和类别清单动手标数据之前我建议你先想清楚三件事每个类别大概有多少样本、图片按什么比例划分训练集和验证集、类别名称怎么统一命名。举个例子我当时的自定义类别有三类defect、scratch、stain所有标注文件里的名字必须严格一致大小写、空格都不能错否则类别映射会乱。类别名称统一用英文小写比如 defect、scratch、stain不要混入标点或空格。训练集和验证集划分比例建议 8:2验证集图片务必和训练集没有交集否则验证结果虚高上线就现原形。样本数量上如果做迁移学习每类至少要有几十到上百张标注图片才比较稳太少的话建议先做二分类或合并相似类别别硬上多分类。目录结构我推荐按这种方式组织后面配置里能直接复用custom_dataset/ ├── train/ │ ├── img_001.jpg │ └── ... ├── val/ │ ├── img_001.jpg │ └── ... └── annotations/ ├── instances_train.json └── instances_val.json2. 环境配置与标注工具准备坑从这里开始2.1 版本对应关系mmdetection 和 mmcv 必须锁死我在这块差点崩溃。mmdetection 2.x 系列必须搭配 mmcv 1.x 系列mmdetection 3.x 则要配 mmcv 2.x这两者不能随意乱配。我当时第一次装直接 pip install mmcv 装了最新版结果 mmdet 2.25 的代码跑起来一堆报错后来才搞清楚版本匹配是硬约束。我这套环境给你作参考Python 3.7 PyTorch 1.10.0 CUDA 11.3 mmcv-full 1.7.0 mmdet 2.25.0。安装方式也比较直接先创建虚拟环境再装 PyTorch然后装 mmcv-full最后装 mmdet。需要提醒的是mmcv-full 的安装会做编译校验你的 CUDA 和 PyTorch 版本不匹配时它会在安装时就报错这时候可以先检查这两个环节。conda create -n mmdet python3.7 -y conda activate mmdet conda install pytorch1.10.0 torchvision0.11.0 cudatoolkit11.3 -c pytorch -c conda-forge pip install mmcv-full1.7.0 pip install mmdet2.25.0版本匹配表在官方文档里写得很清楚但很多人不习惯先查。我的习惯是先确定 mmdet 大版本再回推 mmcv 的对应版本最后根据显卡驱动定 CUDA。多人在同一台服务器上共用环境时建议用 conda 虚拟环境隔离别直接在 base 环境里装容易互相污染。注意如果你跟我一样用 mmdet 3.x 系列配置文件写法会有一点变化比如训练循环相关参数从 Runner 改成了 TrainLoop代码 API 也略有调整。但数据准备和 COCO 格式这条主线是一样的所以本文的核心流程照样适用。2.2 标注工具选型Labelme 比 LabelImg 更适合自定义数据集标注工具有很多我用过 LabelImg也用过 Labelme还试过一些带自动标注功能的新工具。最终我在这个项目里用的是 Labelme原因是它对类别列表的管理更灵活导出 JSON 后转 COCO 格式的脚本也容易写。LabelImg 导出的是 XMLVOC 格式对矩形框标注够用但多边形标注支持一般。Labelme 导出的是 JSON支持矩形框和多边形字段结构清晰适合做后期二次开发。如果目标物体形状不规则建议用 Labelme 的 polygon 模式检测效果反而更好。标注时别急着追求数量先保证质量。我第一次标的时候边框画得随意有的框比物体大了一圈有的只框了一半最后 mAP 卡在 0.4 左右上不去。后来重新检查一遍标注把明显的坏框修掉mAP 立刻涨了五六个点。标注这个环节没有捷径一次做对后面省一堆事。3. 自定义数据集制作从一堆图片到 COCO 格式 JSON3.1 COCO 格式的核心字段你必须弄懂mmdetection 对 COCO 格式的原生支持最好这也是官方示例数据集采用 COCO 的原因。COCO 格式的 JSON 整体是一个大字典最关键的字段是 images、annotations、categories 三个。images 数组里的每个元素至少要包含三个字段id、file_name、width、height。width 和 height 别搞错后面模型会根据这个值来匹配标注坐标。annotations 数组里核心字段是 image_id、category_id、bbox、area、iscrowd。其中 bbox 的格式是 [x, y, width, height]x、y 是左上角坐标width 和 height 是框的宽和高。这里有一个经典的坑很多人习惯用 VOC 的 [x1, y1, x2, y2] 格式直接填进去但 COCO 不认这个它要的是 [x, y, width, height]。如果你转格式的时候不转换训练时目标框位置会乱套模型根本学不出来。area 字段直接填 width * height 就行iscrowd 一般是 0表示单个目标如果一张图里有一堆挤在一起的小目标可以标成 1但新手不建议碰。3.2 用 Labelme 标注的 JSON 转 COCO 的实操脚本Labelme 导出的 JSON 里shapes 字段记录了每个目标的 label 和 points。对于矩形框points 是 [[x1, y1], [x2, y2]] 两个点。转成 COCO 时需要把这两个点换算成 x、y、width、height。我直接把这个脚本放出来这是我在项目里实际用过的版本里面加了基础的容错处理。核心逻辑就是遍历 Labelme 的 JSON 文件读取图片尺寸然后逐条解析标注框最后组装成 COCO 结构。import json import glob import os from PIL import Image def convert_labelme_to_coco(labelme_dir, output_path, categories): images [] annotations [] cat_id_map {name: i 1 for i, name in enumerate(categories)} ann_id 1 image_id 1 for json_path in glob.glob(os.path.join(labelme_dir, *.json)): with open(json_path, r, encodingutf-8) as f: data json.load(f) img_path data[imagePath] img_name os.path.basename(img_path) img_full_path os.path.join(os.path.dirname(json_path), img_name) if not os.path.exists(img_full_path): continue img Image.open(img_full_path) width, height img.size images.append({ id: image_id, file_name: img_name, width: width, height: height }) for shape in data[shapes]: label shape[label] points shape[points] if label not in cat_id_map: continue if len(points) ! 2: continue x1, y1 points[0] x2, y2 points[1] x min(x1, x2) y min(y1, y2) w abs(x2 - x1) h abs(y2 - y1) if w 0 or h 0: continue annotations.append({ id: ann_id, image_id: image_id, category_id: cat_id_map[label], bbox: [x, y, w, h], area: w * h, iscrowd: 0 }) ann_id 1 image_id 1 coco { info: {description: custom dataset}, licenses: [], images: images, annotations: annotations, categories: [ {id: i 1, name: name} for i, name in enumerate(categories) ] } with open(output_path, w, encodingutf-8) as f: json.dump(coco, f, ensure_asciiFalse, indent2) if __name__ __main__: convert_labelme_to_coco( custom_dataset/labelme_jsons, custom_dataset/annotations/instances_train.json, [defect, scratch, stain] )跑完脚本后我推荐你做一次“数据体检”写个简单脚本统计一下 annotation 数量、各类别目标数量、图片分辨率分布。我项目里就出现过某类图片只有十几张、另一类有几百张的严重不平衡如果不提前发现模型会对少数类直接“放弃治疗”。3.3 json 格式检查与常见低级错误转换完成之后要仔细检查几点json 文件编码是不是 utf-8、所有 image_id 是否能在 images 里找到、category_id 是否在 categories 定义范围内。我见过不少人转出来的 jsonannotation 里的 image_id 是从 0 开始的而 images 数组里的 id 是从 1 开始的结果就是训练时一大半标注都被丢弃mmdetection 还会报一堆 warning。还有一个经验图片文件名里不要带中文和空格。虽然 mmdetection 底层能处理但 COCO API 在读取和匹配时偶尔会出问题为了省心前期就直接把所有图片重命名为纯数字加英文字母。如果图片已经在老项目里有固定命名规则建议你在转格式前先批量改名别等训练报错再回头折腾。4. 核心配置文件修改从改着玩到改得明白4.1 解构训练配置文件的关键字典mmdetection 的配置文件本质上是一个 Python 字典文件它会定义模型结构、数据集路径、训练策略等。自定义数据集项目里你基本不需要从零写配置而是从已有的示例配置上改这样少踩很多坑。我以 Faster R-CNN 为例核心的改动点有三个model 里的 num_classes、data 里的 dataset_type 与 classes 和路径、schedule 里的训练轮数和学习率。model 字段里的 roi_head.bbox_head.num_classes 必须改成你的类别数。很多新手第一个报错就是这里num_classes 没改一直用默认的 80加载预训练权重后维度对不上直接 Assertion Error。我这个项目里是 3 类也就是 num_classes3。如果你是单类检测依然要写 1不要写 0。model dict( typeFasterRCNN, backbonedict(typeResNet, depth50), ... roi_headdict( typeStandardRoIHead, bbox_headdict( typeShared2FCBBoxHead, num_classes3, # 改成你的类别数 ... ) ) )4.2 修改数据字典image 路径与 ann 文件缺一不可data 字段是自定义数据集项目的重头戏。你需要显式地告诉 mmdetection训练集和验证集的标注文件在哪里、图片前缀是什么、类别列表是什么。我在实际项目里发现一个容易遗漏的点dataset_type 要写成 CocoDataset同时还要在 train、val、test 三个子配置里分别加上 classes。dataset_type CocoDataset classes (defect, scratch, stain) data dict( samples_per_gpu2, workers_per_gpu2, traindict( typedataset_type, classesclasses, ann_filecustom_dataset/annotations/instances_train.json, img_prefixcustom_dataset/train/ ), valdict( typedataset_type, classesclasses, ann_filecustom_dataset/annotations/instances_val.json, img_prefixcustom_dataset/val/ ), testdict( typedataset_type, classesclasses, ann_filecustom_dataset/annotations/instances_val.json, img_prefixcustom_dataset/val/ ) )samples_per_gpu 表示每张 GPU 卡一次迭代的样本数。如果你的显卡显存只有 8G别硬设成 4通常会 OOM。我一般从 2 开始跑显存有余量再加。workers_per_gpu 是数据加载线程数一般 2 到 4 就够开大了某些老机器 CPU 反而先顶不住。4.3 训练策略配置小数据集别直接 1xmmdetection 的 schedule 配置里有 1x、2x、3x 这些概念分别对应训练 12、24、36 个 epoch。对自定义小数据集我个人不建议直接用 1x 配置因为数据量少模型容易过拟合需要在学习率和训练轮数上做调整。我当时用的是更保守的方案学习率从 0.01 降到 0.001训练轮数延长到 24 个 epoch并配合了余弦退火学习率策略。对小数据集来说低学习率加长训练时间比高学习率短训练的表现稳定得多。当然这也和你的类别数量、标注质量有关我的经验是损失函数在前 10 个 epoch 下降明显后面趋于平缓所以 24 个 epoch 基本是性价比最高的。lr_config dict( policyCosineAnnealing, min_lr_ratio0.001 ) runner dict(typeEpochBasedRunner, max_epochs24)提示如果你用的是 mmdet 3.xrunner 配置换成了 train_cfg 里的 max_epochs写法上略有差别但调整思路完全一致。5. 训练、评估与可视化跑起来只是第一步5.1 启动训练的命令与日志解读环境没问题、配置改完之后就可以启动训练了。单卡训练的命令很简单python tools/train.py configs/my_config/faster_rcnn_r50_fpn_1x_custom.py如果你想在后台跑并且把日志写到文件里就用nohup python tools/train.py configs/my_config/faster_rcnn_r50_fpn_1x_custom.py --work-dir work_dirs/my_custom train.log 21 训练刚开始时日志里会出现模型结构预览、数据集加载信息然后是每个 iter 和 epoch 的 loss 数值。我一般会重点关注 loss_rpn_cls、loss_rpn_bbox、loss_cls、loss_bbox 这几个指标。正常情况是前几个 epoch 整体 loss 明显下降后面缓慢下降。如果 loss 直接变成 NaN大概率是学习率过大或者数据里有异常标注先查这两个。多卡训练我也提一句mmdetection 提供了分布式训练脚本两卡就写 2四卡就写 4bash tools/dist_train.sh configs/my_config/faster_rcnn_r50_fpn_1x_custom.py 2但多卡训练有一个隐藏大坑workers_per_gpu 的数据加载如果不同步某些机器会把显存占满然后报共享内存错误这时可以把 DataLoader 的 persistent_workers 和 num_workers 调低。5.2 评估模型用 mAP 和损失曲线说话训练完成后用 test 脚本评估模型在验证集上的表现python tools/test.py configs/my_config/faster_rcnn_r50_fpn_1x_custom.py work_dirs/my_custom/epoch_24.pth --eval bbox输出的 mAP 是综合评价指标iou0.5 的 AP 对实际项目最有参考价值。如果 AP 在 0.5 以下先别急着调网络结构优先检查数据问题比如标注框是否准确、类别是否平衡、train 和 val 是否有交集。我那个项目里第一次评估 mAP 只有 0.42排查之后写上一个小脚本统计每个类别的 AP发现 scratch 类别的 AP 比其他类低了近 20 个点原因是这个类别的样本只有二十几张增强也不够后续加数据和做数据增强后AP 提到了 0.68。评估结果会输出一个包括 mAP、mAP0.5、mAP0.75 的表格。如果你要部署到实际业务里建议把重点放在 mAP0.5 和对应类别的 AP 上这两个指标更贴近业务真实感受。mAP0.75 偏严格适合对定位精度要求高的场景。5.3 推理可视化看一眼模型到底检测成了什么样评估指标只是一个数字真正确认模型效果的方式还是可视化推理。mmdetection 2.x 自带 demo 脚本可以直接对单张图片做推理python demo/image_demo.py custom_dataset/val/test_img.jpg configs/my_config/faster_rcnn_r50_fpn_1x_custom.py work_dirs/my_custom/epoch_24.pth --device cuda:0 --score-thr 0.3我实际用过之后觉得可视化环节不只是“看效果”更是发现数据问题的机会。比如模型把很多背景误检成 defect那可能是训练集里负样本太少图片背景过于单一如果漏检很多目标本身就很小的物体可能需要提高输入图像分辨率或者调整 anchor 的 scale。这些判断如果只看 mAP 数字容易一头雾水。写推理脚本也不难核心是加载 config 和 checkpoint构建模型然后调用 inference_detector。这块代码网上很多但在 mmdet 2.x 和 3.x 之间 API 有差异建议直接翻对应版本的官方 demo 文件改别跨版本硬套。6. 常见问题与排查技巧实录6.1 高频报错速查表我整理了一份常见报错速查表每一行都是我在实践里遇到过的真实问题你可以直接对照排查。报错/表现原因解决方案AssertionError: num_classes 不匹配配置里忘了改 ROI head 的类别数改 model.roi_head.bbox_head.num_classesKeyError: xxx is not in models registrymmdet 版本和配置文件版本不一致确认 mmdet 2.x 配 mmcv 1.x3.x 配 mmcv 2.xOOM / CUDA out of memorybatch size 或输入分辨率过大调小 samples_per_gpu关掉 albu 增强IndexError 或 KeyError 在加载数据集COCO JSON 里 image_id 或 category_id 对不上写脚本校验 annotation 和 images 的对应关系loss 变 NaN学习率过大或标注里出现异常的框调低 lr检查标注 bbox 是否为 0 或负数验证 mAP 太低但在训练集很好训练集验证集有交叠或严重过拟合重新划分数据集减少 epoch加数据增强图片加载时报文件不存在img_prefix 路径不对或图片名带中文/空格统一重命名图片核对目录结构6.2 训练过程中几个值得留意的细节第一个细节是 checkpoint 的保存策略。mmdetection 默认每 1 个 epoch 存一个 checkpoint如果训练 24 个 epoch会占不少磁盘。我在项目里一般设置每个 3 个 epoch 间隔保存一次同时保留 best 模型既不影响后期挑模型也不至于把磁盘撑爆。第二个细节是数据增强。mmdetection 默认的增强主要是随机翻转和缩放对小数据集来说略保守。我后来在 train_pipeline 里加了多尺度训练增强效果提升明显。但要注意增强不是越多越好有些增强会把小目标直接裁掉反而降低性能建议逐步加每次改动后重新评估对比。第三个细节是预训练权重的选择。mmdetection 会自动下载在 ImageNet 上预训练的 backbone 权重如果你的网络环境差下载会卡住最后训练半天还在等权重。提前把权重文件下载下来放到对应目录下然后配置里指定 pretrained 路径可以省很多时间。6.3 我的独家避坑技巧先跑通十张图再全量训练这个方法我想推荐给所有人。正式启动全量训练前先做一个 mini 实验只拿 10 张训练图片和 3 张验证图片把完整的训练、评估、推理流程跑一遍。如果这个小实验能顺利跑完说明环境、配置、数据格式都没问题如果这个过程里报错你只需要盯着十张图调试问题定位快得多。我早期没这个习惯直接 2000 张图开跑结果半小时后才发现 COCO JSON 里 class 顺序写反了白白浪费一大笔时间。后来每次新建数据集项目我都是先跑十张图确认流程顺畅后再切到全量数据。这个习惯帮我省下的时间真的够写好几篇博客了。6.4 关于“验证集效果不错但实际场景翻车”的一点思考这个问题不是个别现象。模型在验证集 mAP 很高一到实际拍摄的新图片上效果就很差最常见的原因是验证集和实际测试数据分布不一致。实际拍摄的图片可能光照更复杂、角度更多变、目标更小而这些在你的标注图片里根本没见过模型自然泛化不掉。我给自己的项目加了一个额外的小测试集从真实场景里挑了一部分图片单独做一次评估不参与训练也不参与常规验证。每次迭代模型后都拿这个测试集跑一遍对比 mAP 变化。这个做法不复杂但能帮你更早发现模型的真实短板。我在实际项目中还有一个体会train 和 val 的划分不能简单随机打散尤其是同一场景拍出来的连续帧图片几乎一样的背景如果一部分进了训练集另一部分进了验证集评估结果会虚高不少。我后来改成按视频片段或者按拍摄时间切分数据尽量保证同场景的图片不会被同时分到两边。写在最后一点真实感受整条自定义数据集流程走下来我最深的感受是mmdetection 把模型的训练门槛降得很低但数据工程的门槛一点没降。很多人一上来就折腾配置、折腾模型结构其实真正决定项目成败的往往是你有没有把数据转好、把类分对、把训练验证集切得干净。我这几个星期踩过的坑差不多都写在前面了。如果这个项目让你重新走一遍我建议你直接从“十张图 mini 实验”开始先确认环境通、数据通再全量开跑真的能少走很多弯路。之后再遇到自定义数据集的检测需求就不会像我第一次那样手足无措了。