开头写了,然后按章节展开。我尽量让内容有真实博主的经验感,有具体数字和细节,不空泛。 “跑通了官方demo,把大象、斑马、汽车的掩码看得明明白白,结果一换上自己的数据集,loss像心电图,预测图上一片马赛克。”——这个场景我猜你并不陌生。Mask R-CNN(实例分割模型)在COCO上的效果确实惊艳,但实际工作中最常遇到的问题,恰恰是“如何用自己的数据训练它”。这篇文章就是围绕“使用Mask R-CNN训练自己的数据集”这个主题,把数据准备、格式转换、训练参数、损失分析、推理导出全流程拆开讲。适合已经跑通官方sample、想把自己的业务图像做成实例分割模型的开发者,也适合刚接触深度分割任务、想系统了解训练链路的学生。
我最早用这个模型是在一个零件分拣项目上,需要把传送带上重叠摆放的工件逐个分割出来,再计算每个工件的轮廓面积。当时踩了很多坑,后来回头看,问题大多集中在三个地方:数据标注不规范、超参照搬COCO、以及完全不看训练日志。下面把我反复验证过的流程和心得完整写出来。
1. 内容整体设计与思路拆解
1.1 为什么选择Mask R-CNN而不是纯目标检测模型
先明确一个问题:如果只需要框住目标,用检测网络就够了,为什么要上Mask R-CNN?因为很多真实场景里,“知道物体在哪”不够,还要知道“物体长什么样”。
拿我的零件分拣项目举例,两类工件外观接近但轮廓差异很大,检测框会重叠,单靠矩形框根本无法区分谁是谁。Mask R-CNN在检测框的基础上多了一个mask分支,输出像素级的分割掩码,所以能同时解决三个问题:目标在哪、目标是什么、目标占据哪些像素。这对遮挡场景特别关键——两个物体叠在一起,检测框只能给出一个大致区域,而mask能把每个物体的像素边界分开,后续对mask求连通域,还能精确计算面积、几何中心、姿态角。
从结构上看,Mask R-CNN是在目标检测模型基础上加了“分割分支”的双阶段架构:主干网络提取特征,RPN(区域建议网络)生成候选框,RoIAlign完成像素对齐,最后检测头和分割头并行输出。它的核心创新是RoIAlign,解决了之前RoIPooling在像素级对齐上的精度损失问题。对自定义数据集而言,这意味着必须手动调整的“坑点”主要在数据格式和训练配置,而不是网络结构本身。
1.2 迁移学习与Backbone选型怎么定
我的建议很直接:训练自己的数据集,除非你有十万级以上的海量数据,否则一定用COCO预训练权重做迁移学习,不要从头训练。COCO权重在特征提取层面已经学到了丰富的通用视觉特征——边缘、纹理、局部形状,这些对多数业务场景都够用。我们需要做的,是让模型“适应”新类别的语义,而不是教它重新认识世界。
具体实施时,常用的Backbone有ResNet50和ResNet101。两者的差别是:ResNet50更快、显存占用更低,适合小数据集和快速迭代;ResNet101特征更丰富,精度通常高一点,但对显存要求高,训练时间也几乎翻倍。我做过对比实验,在大概1000张自定义图像上,ResNet50只比ResNet101低了不到2个百分点的AP50,但训练速度却快了近一半,这显然适合快速跑实验。
另外要特别注意一个误区:不少人一上来就“冻结”Backbone只训练后几层。在小数据集上这有一定道理,但对Mask R-CNN这种多任务模型,建议Backbone部分微调而不是完全冻结。因为你的数据分布和COCO差异很大时,深层特征也要适应新域,而“冻结”会让分割分支学得很慢。我通常的做法是前几个epoch保持Backbone的学习率极低,后面再放开。
1.3 全流程拆解:从图像到能用的分割模型
整个项目我习惯分为五个环节,缺一个都会在后期暴露问题:
- 数据准备与标注:采集图像、筛选、清洗,同时完成标注格式定义。
- 格式转换与数据集划分:把标注结果转成COCO格式,划分训练/验证集,写类别元数据。
- 训练配置与调参:修改类别数、训练超参,启动训练并实时监控损失。
- 损失分析与迭代优化:理解每一路损失的意义,判断模型是欠拟合还是过拟合。
- 推理验证与模型导出:在测试图上可视化预测结果,评估指标,确定可以使用的格式。
每个环节的具体操作下面单独展开,我尽量写清楚“为什么这么做”和“我遇到过的反面教材”。
2. 数据集准备:标注规范和格式转换
2.1 图像采集与筛选:别急着标,先做减法
图像采集的阶段最容易被人忽略。很多人随便找几百张图就开始画多边形,结果画到一半发现数据全是同一场景同一角度的,“换个角度模型就废了”。我自己的经验是,采集阶段就做好三件事:
- 覆盖多样性:至少包含不同角度、不同光照、不同背景、不同距离的图像。如果数据来自视频抽帧,必须先做去重,否则相邻帧几乎一样,训练集和验证集会高度相似,评估结果虚高。
- 数量底线:单类目标最少也要50~100张起了才能看出趋势。这不是说几十张不能训,而是效果不稳定,验证集的随机波动会很大,无法判断调参方向是否正确。
- 前景占比把控:目标在图像里太小,对分割模型非常不友好。如果一个物体只占图像1%的像素,RPN很难生成高质量的候选框。如果项目场景里目标就是小,可以在标注时配合切图策略,把大图切成小块再标,提高目标的相对面积。
2.2 标注工具的选取与多边形标注规范
标注工具我用过好几款,最常用的是Labelme。它的好处是标注结果是一个很简单的JSON,包含多边形的顶点坐标列表和标签名,后续转COCO格式很方便。对于什么类要标、怎么标,必须提前定一套自己的规范,否则不同人标出来的数据模型学出来的东西会很怪。
我踩过的一个典型坑是:多人标注时,“目标被遮挡”的处理标准不一致。有人把被挡住的边界也描出来,有人只描可见部分。Mask R-CNN训练时mask是像素级的,这两种标注混在一起会让分割分支无所适从。
所以我们内部定的标注规范是:
- 目标被遮挡超过50%时,不标注该目标,只标注完整可见的目标;
- 目标局部被遮挡但主体轮廓可辨认时,只描可见区域,不要脑补被挡住的边界;
- 两个紧贴的目标,多边形之间至少留1~2像素空隙,避免掩码边界重叠;
- 背景误识别为目标的都要标出来,特别是那些“容易混淆但又不属于目标类别”的物体,要么归为背景,要么单独建类。这一点对实例分割特别重要,如果模型把背景错当目标,多半是标注时漏标了相似物。
标注完成后,还要做一次质量检查,把每个类别的实例数量统计一遍,看有没有标注遗漏。这一步虽然耗时,但比训练完发现模型“幻觉”再回头补标要省事得多。
2.3 Labelme JSON转COCO格式:一个脚本解决核心问题
Mask R-CNN对外公开的实现里,数据加载通常按COCO格式来组织。COCO数据集的标注格式是一个大JSON,里面分images、annotations、categories三个部分,而Labelme是每张图一个JSON。所以核心工作是写一个转换脚本,把这些零散的标注文件汇总成COCO格式。
转换脚本的核心逻辑很简单:
import json, os, glob from collections import defaultdict def labelme_to_coco(img_dir, json_dir, out_json, categories): # categories示例: [{"id": 1, "name": "object_1", "supercategory": "object"}] coco = { "images": [], "annotations": [], "categories": categories, } ann_id = 1 for img_id, json_path in enumerate(sorted(glob.glob(os.path.join(json_dir, "*.json"))), 1): with open(json_path, "r", encoding="utf-8") as f: data = json.load(f) img_info = { "id": img_id, "file_name": data["imagePath"], "width": data["imageWidth"], "height": data["imageHeight"], } coco["images"].append(img_info) for shape in data["shapes"]: # 将多边形顶点拆成 COCO 的扁平列表 [x1,y1,x2,y2,...] points = [coord for pt in shape["points"] for coord in pt] # 这里还需要计算bbox,COCO的bbox格式为 [x, y, w, h] x_coords = [pt[0] for pt in shape["points"]] y_coords = [pt[1] for pt in shape["points"]] x_min, y_min = min(x_coords), min(y_coords) w, h = max(x_coords) - x_min, max(y_coords) - y_min # segmentation 里可以填多边形点集;对于多边形标注,isbbox 建议设为0 ann = { "id": ann_id, "image_id": img_id, "category_id": ..., "bbox": [x_min, y_min, w, h], "segmentation": [points], "area": w * h, "iscrowd": 0, } coco["annotations"].append(ann) ann_id += 1 with open(out_json, "w", encoding="utf-8") as f: json.dump(coco, f, ensure_ascii=False, indent=2)这里有两个细节:bbox很多人直接用多边形最小外接矩形,公式是对的,但计算时不要用int去截断,建议保留浮点精度再四舍五入,因为小目标的对象,几个像素的偏差都会直接影响分割结果的对齐。另一个是area字段,建议也用多边形真实面积而不是w*h,COCO评估里计算mask AP时对面积分组很敏感。
数据划分上,建议图像按8:2划分训练集和验证集。注意划分时按“场景”划,不要让同一个物体的不同帧同时出现在训练集和验证集里,否则验证集等于开卷考试。
2.4 数据增强:决定模型泛化能力的隐藏因素
Mask R-CNN官方实现里内置了几种数据增强:随机水平翻转、随机亮度和对比度调整、图像缩放。这几项对小数据集已经很有帮助,但还不够。我在实践中补充了随机旋转(范围±15度),以及随机裁剪,效果都很明显。
需要明确的是,数据增强不是越多越好。实例分割的mask是像素级标签,有些增强操作会破坏标签的几何一致性。比如旋转时mask和多边形顶点要一起旋转;随机裁剪时要保证裁剪区域内没有“截断的目标”,或者把截断的目标从标注里同步删掉。如果简单地对图像做裁剪但标签没跟住,模型就会学到“物体的上半身是一个类”。
类别不均衡的问题也要在这里处理。如果你要分割的目标有三类,A类200张、B类50张、C类20张,B和C就是短板。我在项目里的做法是:对少数类图像做复制粘贴增强,先把它们的样本数拉到一个基线(比如都到80张以上),再统一做常规增强。要记住,实例分割模型对长尾类别的学习能力远比你想象中弱,类别数一旦不均衡,loss会被大类主导,小类几乎学不出来。
3. 训练配置与实操:从参数到命令
3.1 环境准备与源码选择
环境部分我直接给一套经过验证的组合:深度学习的GPU显存建议至少8GB,能用11GB以上最好。训练框架用TensorFlow 2.x + Keras;COCO分割任务的标准开源实现,社区里最常用的就是那个带“mrcnn”模块的代码库。
注意一点:拿到源码后不要急着跑,先检查依赖里是否有pycocotools。这个库是COCO官方评估工具包的编译版,很多人在Windows上装它吃苦头。如果编译安装失败,优先找预编译的whl包,实在不行就退一步,先把训练跑通,评估指标暂时只看loss。不要因为一个小工具卡住整个项目。
3.2 修改配置文件:核心是类别数和数据集路径
官方实现里,配置类通常在mrcnn/config.py里面,你之前跑demo时可能已经接触过。训练自己的数据集,需要新建一个配置类继承自原来的基础配置,重点覆盖下面几项:
class CustomConfig(Config): # 类别数,包含背景类。比如要分割3类目标,这里就填4 NUM_CLASSES = 1 + 3 # 训练/验证图像的尺寸,按你的图像实际大小来 IMAGE_MIN_DIM = 800 IMAGE_MAX_DIM = 1280 # 小数据集建议开小步长,每轮迭代少一点 STEPS_PER_EPOCH = 200 VALIDATION_STEPS = 50 # 训练分支:heads / all # heads表示只训练网络头部,all表示微调全部层 LEARNING_RATE = 0.001NUM_CLASSES是最大的坑点之一。它必须等于“你的类别数+1”,加的是背景类。我第一次训练时直接填了3,结果模型在加载预训练权重时维度对不上,报了一堆错。如果你是二进制分割(只区分目标和背景),这个值就是2。
IMAGE_MIN_DIM和IMAGE_MAX_DIM也值得多说一句。官方默认是1024,但如果你的图像分辨率很高,比如工业相机拍的3000×4000图片,直接喂进去要么显存炸掉,要么被压得面目全非。我的建议是不要一味追求大分辨率,先缩放到800~1280区间,让模型跑通。等效果稳定了,再专门针对小目标提升输入分辨率。
3.3 训练超参的选择:别照搬官方默认值
超参数这部分,我把自己试出来的实际经验整理成一张表:
| 超参数 | 官方默认 | 自定义小数据集建议 | 理由 |
|---|---|---|---|
| LEARNING_RATE | 0.02 | 0.001~0.005 | 大数据集可以用大学习率,小数据集用大学习率容易发散 |
| WEIGHT_DECAY | 0.0001 | 0.0001 | 正则可保留,但不要加太猛,否则小模型拟合不足 |
| BATCH_SIZE | 2 | 2~4 | 显存允许下尽量4,BN在batch太小时不太稳定 |
| EPOCHS | 30+ | 20~40 | 视loss下降情况而定,多注意验证集指标 |
| RPN_ANCHOR_SCALES | (32, 64, 128, 256, 512) | 根据目标尺寸调整 | 目标小就把anchor调小,目标大就调大 |
关于RPN_ANCHOR_SCALES,这是新手最容易忽略的:anchor是RPN生成候选框的“尺子”。如果你的目标在图像里普遍小于64×64像素,默认的anchor根本匹配不上;而如果目标是很细长的棒状物,正方形anchor效果也不好。一个经验做法是:统计你标注数据的bbox宽高分布,然后用这些统计值去设置anchor。比如我的工件数据里,最小框约30×40,最大框约200×150,我就把anchor设成(32, 64, 128, 256),并在训练前用一小段代码做anchor匹配率检查。
训练后的启动命令,我通常是这样跑的:
python custom/train.py \ --dataset /path/to/your/dataset \ --subset=train \ --model=coco--model=coco表示从COCO预训练权重开始,--model=last表示从最近一次的训练中断点继续。中途断了不要慌,直接用--model=last接着跑就行,这是官方实现里比较方便的地方。
3.4 训练过程中的日志监控
开始训练后,不要干等着。训练日志里每一行都会打印loss、rpn_class_loss、rpn_bbox_loss、mrcnn_class_loss、mrcnn_bbox_loss、mrcnn_mask_loss这几个值。我要强调一个心态:前几个epoch看到loss在2.0甚至3.0以上是很正常的,尤其是mask_loss刚从随机状态开始学习时。如果连续十几个epoch一点下降趋势都没有,才是真的要找问题。
真正值得警惕的信号是:训练集loss持续下降但验证集指标恶化,那是过拟合;训练集loss一直不动在某个高位,那是学习率太小或数据有问题;loss先降后升且值突然无限大,那是梯度爆炸。后两种在自定义数据集里的出现频率其实比过拟合高得多,多半是标签里出现了NaN值,或者某个类别的标注极度不规范。
我习惯在训练时每隔一段时间就把验证集的预测结果可视化保存一批,用肉眼看比看一百行数值统计更快发现模型问题。官方实现里带了一个inspect_model的脚本可以干这个事,后面章节会细说。
4. 损失曲线分析与模型迭代调优
4.1 五种损失分别代表什么
训练日志里五种损失很容易让人一头雾水,我逐个解释:
- rpn_class_loss:RPN生成的候选框“是否包含目标”的分类损失。如果这个值一直高,说明候选框的生成质量差,正负样本失衡严重。
- rpn_bbox_loss:候选框位置回归损失。反映RPN建议框的位置准不准。
- mrcnn_class_loss:对候选框进行目标类别分类的损失。注意,“背景被错分为目标”会体现在这里。
- mrcnn_bbox_loss:最终检测框的回归损失。
- mrcnn_mask_loss:掩码分割的损失,计算每个像素预测掩码和真实掩码的交叉熵。
其中mrcnn_mask_loss是这个模型独有的。它的值通常比前两个大,因为像素级分类的样本量远大于框级的样本量。我见过不少人在日志里看到mask_loss在1.0以上就慌了,其实那完全正常。判断标准应该是它是否在稳定下降。只要同一epoch内loss比上一个epoch低,就是健康的。
4.2 用损失曲线判断训练是否健康
为了更直观地判断,我后来都会开启TensorBoard可视化训练过程,把每个epoch的多个损失曲线放在同一个面板里看。具体判断方法:
- 如果
rpn_bbox_loss降得很快但mrcnn_mask_loss很慢,说明主干特征还没充分适配你的数据,可以考虑提高backbone部分的学习率倍率,或者延长训练epoch。 - 如果
mrcnn_class_loss很高而mrcnn_mask_loss不高,说明模型能分像素但认不出类别。这种情况通常是类别样本严重不均衡,或者训练集里类别标注错乱。 - 如果所有损失都在前10个epoch内快速下降、之后完全平了,而且验证集上的可视化结果已经开始出现不错的mask轮廓,说明模型已经接近当前数据规模下的“天花板”,继续硬训只会过拟合。
我最常犯的错是“无脑训满30个epoch”。后来养成了每5个epoch保存一次模型、跑一次验证集可视化的小习惯,看到mask基本能把目标轮廓描下来,就真实停手。深度学习工程讲究“快速验证、小步迭代”,模型训练要的是达到业务可用阈值,不是让训练曲线彻底归零。
4.3 调优空间:从anchor到NMS再到后处理
当模型能跑通但精度不满意时,按优先级往下调:
- 调整anchor尺度。先统计标注数据的宽高分布,按聚类结果改
RPN_ANCHOR_SCALES。这一项对于目标尺寸和默认值差别大的数据集,收益是立竿见影的。 - 调节RPN NMS阈值。官方默认
RPN_NMS_THRESHOLD是0.7,如果目标密集,建议降到0.5,减少重叠候选框对结果的干扰。 - 调整检测NMS阈值。
DETECTION_NMS_THRESHOLD控制最终输出框的去重。目标相互遮挡严重时,0.3会比0.5更合适,但也会让一些漏检找回来,需要在漏检和误检之间权衡。 - 后处理:mask输出后经常会有些小空洞和毛刺,可以用形态学开闭运算处理。在工业项目上,这一步虽然笨,但对下游的尺寸测量很关键。我的工件测量里就用了一次闭运算加一次开运算,面积计算的结果稳定了很多。
4.4 小目标分割的专项优化
如果你的任务主要面向小目标,上面这些常规调优还不够。我建议做两件事:
- 提高输入分辨率。把
IMAGE_MAX_DIM从1024提到1536甚至2048。代价是显存占用飙升、训练速度变慢,所以在显存允许的前提下做,而且要结合数据增强里的随机裁剪。 - 在标注时做切图。把原图切成若干有重叠的patch,每个patch单独标、单独训练。这样能把小目标的相对尺寸放大,减少RPN漏检。代价是要处理patch拼接时的重复预测,推理时需要对重叠区域做NMS合并。这个方案实现起来有点麻烦,但对小目标数据集几乎是效果最明显的手段。
5. 推理验证与模型导出的完整路径
5.1 可视化预测:训练完先别急着算指标
训练结束后,先别用脚本算一堆mAP,先跑一张验证集上的预测图,直观看看分割效果。官方仓库里的inspect_weights和inspect_model脚本能加载训练好的权重并可视化mask和roi。
我的做法是写一个独立脚本,加载模型,然后用model.detect()对验证集图像做预测,把rois、class_ids、masks三个输出叠加在原图上保存。这里有一个重点:model.detect()返回的坐标是经过resize后的坐标,回画原图时要按image_meta里的原始尺寸换算。很多第一次接触的人在这里踩坑,画出来的mask位置全部偏移。
关于预测阈值,官方默认检测置信度阈值是0.7,对小数据集建议先降到0.5看效果。这个指标没什么高深理论,完全取决于你的业务对“漏检”和“误检”谁的容忍度更低——如果分割错了会带来严重后果,阈值就调高;反之可以调低。
5.2 用COCO评估指标量化效果
可视化的感觉对了,再用评估脚本算出量化指标。这里推荐关注三个指标组合:AP50、AR@10和AP_small(如果小目标多的话)。不要只看一个mAP。
比如我的工件项目,最终模型AP50到了0.86,但AP_small只有0.42。这说明模型对大目标很稳,小目标还是吃亏。如果业务里小目标很常见,就必须回到4.3的优化方案里继续做;如果业务里小目标不重要,那这个模型就可以上线了。做决策的是业务需求,而不是指标本身。
5.3 模型保存与部署的可能路径
训练完的权重是Keras的h5格式,但这不直接等于“能上线”。我常用的导出方式有两种:
- 服务端部署:直接把h5模型加载到TensorFlow Serving或者基于Flask封装成HTTP接口。因为Mask R-CNN是两三层的模型结构,服务端推理没有特别大的瓶颈,只要GPU内存够。
- 移动端或边缘设备部署:需要把模型转成TensorFlow Lite或ONNX格式。转换时首先要把输入图像预处理(减均值、除方差、resize)提前到前端完成,后端模型只做推理。要注意,Mask R-CNN因为包含RPN、RoIAlign、mask解码等结构,直接转ONNX经常会碰到自定义算子的支持问题。我在边缘设备上试过一次,最终是把模型拆成了“特征提取+检测头”和“分割头”两个子图分开转,才跑通了整个流程。
无论哪种部署方式,都建议先在推理脚本里明确固定一个min_confidence和nms_threshold,并记录到模型的配置文件中,避免开发和部署环境两套参数不一致。
6. 常见问题与排查技巧实录
最后这部分,我把这几年来被问得最多、也亲身趟过的一些问题整理成速查表,按照训练流程的先后顺序排好。
| 现象 | 可能原因 | 解决思路 |
|---|---|---|
| 加载COCO JSON时报错,提示找不到图片 | 图像数量与JSON里images记录不一致 | 检查标注时是否乱动过图片文件名;用脚本重新统计JSON里的文件名,和实际文件夹做差集 |
| 训练几轮后loss变成NaN | 标注坐标出现无穷值或bbox为空 | 在转换脚本里加校验,过滤掉顶点数小于3、面积小于0或坐标超界的标注 |
| GPU显存不足 | 输入分辨率过大或batch size太大 | 先降到batch size=2,再考虑下调IMAGE_MAX_DIM;虚拟显存技术兜底但会拖慢训练 |
| loss下降缓慢且始终在3.0以上 | 学习率太小,或backbone冻结太死 | 调大LEARNING_RATE到0.005;关闭冻结,让backbone最后一阶段也参与微调 |
| 训练集表现好,验证集预测出现大量背景误检 | 标注时漏标了相似背景物体 | 回到数据集,补齐漏标,不要急着调参 |
| 模型预测的mask位置正确但形状像“膨胀”了 | 数据增强里的旋转/裁剪没有同步修改mask标签 | 检查增强代码,确保对图像和标注做相同的几何变换 |
| 小目标完全检测不到 | anchor尺度不匹配小目标 | 统计目标bbox尺寸,把anchor改小或增加更小的尺度 |
| 切换验证集评估分数忽高忽低 | 验证集划分有泄漏,或验证图像太少 | 按场景划分而非按单图划分,增加验证集数量 |
再补几个容易踩的暗坑:
- 多GPU训练时,如果代码里写的batch size是2,实际批大小要乘以GPU数量。学习率也要按比例放大,否则梯度更新步长和正常单卡不一样。
- 不要直接把COCO预训练权重的类别数改成你的类别数后强行加载。加载前需要把最后几层的权重剔除,否则维度不匹配会报错;用
load_weights(path, by_name=True)配一个exclude列表是稳妥做法。 - 小数据集别用太大的随机失活率(dropout)系数。数据量本来就少,过强的正则会直接让模型欠拟合,表现为训练集loss都降不动。
我这个项目后来还扩展了一下:把Mask R-CNN输出的mask拿去计算每个工件的像素面积,再结合标定换算成真实尺寸,做成了流水线上的尺寸分拣逻辑。表面上看这只是“分割模型的后续应用”,实际上它提示了一个很重要的事——Mask R-CNN常常不是终点,而是视觉系统的中间层输出。训练好一份自己的数据集,换来的不仅仅是一个模型,更是一整条可以反复利用的图像理解流水线。
如果让我给一个最后的建议,那就是:首轮训练一定要做“极简闭环”。数据先选50张、类别先只跑一类、训练先只跑5个epoch,把所有格式和流程问题全部引爆,再逐步扩充。这样走一遍,远比一开始就标500张、训30个epoch高效得多。这个习惯我沿用至今,几乎每一次自定义训练都能靠它保住大局。