从零开始:使用Mask2Former训练你的专属数据集
1. 环境搭建为Mask2Former铺好路想玩转Mask2Former第一步就是把它的“家”给搭好。这个“家”就是你的开发环境。很多新手朋友一上来就被各种依赖和编译搞得头大其实跟着步骤走踩几个我当年踩过的坑你也能轻松搞定。记住环境配置的核心就两点版本对齐和耐心。首先你得把Mask2Former的源码“请”到你的电脑里。直接从GitHub上克隆是最稳妥的办法。打开你的终端找个你喜欢的目录执行git clone https://github.com/facebookresearch/Mask2Former.git。这一步基本不会出错网速快慢而已。源码下载下来后先别急着往里冲我们得看看它的“左膀右臂”——Detectron2。Mask2Former是构建在Facebook的Detectron2框架之上的所以Detectron2是必须的依赖。官方推荐的做法是把Detectron2的源码也克隆下来放到Mask2Former的根目录里。你可以执行git clone https://github.com/facebookresearch/detectron2.git然后把整个detectron2文件夹拖进Mask2Former文件夹里。这样它们俩就在一个屋檐下了。接下来就是安装环节。这里有个小坑直接pip install -e .安装Mask2Former可能会因为PyTorch或CUDA版本不匹配而报错。我的经验是先别管Mask2Former确保你的PyTorch是带CUDA版本的并且版本不要太新也不要太旧PyTorch 1.10到1.12配合CUDA 11.3或11.6是比较稳妥的组合。确认无误后我们先进入Detectron2目录进行编译安装。有时候需要指定CUDA架构比如你的显卡是RTX 30系算力8.6但为了兼容性可以尝试用TORCH_CUDA_ARCH_LIST8.6 FORCE_CUDA1 python setup.py build develop。不过对于大多数教程提及的通用情况在Mask2Former的目录下直接运行pip install -e .通常就能自动处理好Detectron2的依赖。如果报错再回头单独处理Detectron2。安装完主体框架还需要两个重要的工具库panopticapi和cityscapesScripts。它们是处理特定数据集格式用的即便你只用自己的简单数据集安装上也能避免一些潜在的导入错误。命令很简单pip install githttps://github.com/cocodataset/panopticapi.git和pip install githttps://github.com/mcordts/cityscapesScripts.git。最后别忘了安装项目要求的其他Python包执行pip install -r requirements.txt。到这里软件层面的环境就差不多了。但还没完Mask2Former模型里用到了一个为了加速而写的CUDA算子我们需要编译它。这个步骤是很多新手失败的地方。你需要进入mask2former/modeling/pixel_decoder/ops/这个目录然后运行里面的make.sh脚本。在运行前我强烈建议你先打开这个make.sh文件看一眼。里面通常就是一行python setup.py build install。有时候直接运行sh make.sh会失败可能是因为环境变量问题。我踩过的坑是手动在终端里设置一下CUDA_HOME路径比如export CUDA_HOME/usr/local/cuda-11.6请替换成你的CUDA安装路径然后再执行python setup.py build install成功率会高很多。看到编译成功生成.so文件这一步才算真正过关。2. 数据准备把你的图片变成模型认识的语言环境好了接下来就是喂给模型“食物”——数据。Mask2Former这类模型通常吃的是COCO格式的数据。但我们的原始数据往往是每张图片配一个LabelMe生成的json标注文件这就像你说中文模型只懂英文我们需要一个“翻译”过程。首先在项目根目录下建立一个清晰的数据文件夹比如就叫datas。里面再创建两个子文件夹train和val分别存放训练集和验证集的图片及对应的LabelMe JSON文件。你的数据结构应该长得像这样datas/ ├── train/ │ ├── 1.jpg │ ├── 1.json │ ├── 2.jpg │ ├── 2.json │ └── ... └── val/ ├── 100.jpg ├── 100.json └── ...注意图片格式jpg、png都行但一定要和json文件的主文件名一一对应。这个看似简单的步骤却是后续所有工作的基础务必检查清楚。现在到了关键的一步格式转换。我们需要写一个小脚本或者利用现成的工具把这一堆LabelMe的JSON文件合并转换成两个COCO格式的JSON文件instances_train.json和instances_val.json。网上有很多labelme2coco.py的脚本你可以找一个靠谱的。这里我分享一个我常用的脚本核心思路你需要根据实际情况调整。这个脚本的核心是调用labelme2coco库如果你安装了labelme它可能自带。假设你的脚本和datas文件夹在同一级目录脚本内容大致如下import os import json import argparse from labelme import utils import numpy as np from PIL import Image from tqdm import tqdm def labelme_to_coco(labelme_dir, output_json_path, category_list): labelme_dir: 存放图片和json的文件夹如datas/train output_json_path: 输出的COCO格式json路径如datas/annotations/instances_train.json category_list: 类别名字列表如 [cat, dog] # 创建COCO JSON的基本结构 coco_output { info: {}, licenses: [], categories: [], images: [], annotations: [] } # 构建类别字典 for i, cat_name in enumerate(category_list, start1): coco_output[categories].append({ id: i, name: cat_name, supercategory: none }) category_name_to_id {cat[name]: cat[id] for cat in coco_output[categories]} # 遍历文件夹 image_id 1 annotation_id 1 label_files [f for f in os.listdir(labelme_dir) if f.endswith(.json)] for label_file in tqdm(label_files): label_path os.path.join(labelme_dir, label_file) with open(label_path, r, encodingutf-8) as f: label_data json.load(f) # 处理图片信息 image_file_name label_data[imagePath] image_path os.path.join(labelme_dir, image_file_name) img Image.open(image_path) width, height img.size coco_output[images].append({ id: image_id, file_name: image_file_name, height: height, width: width }) # 处理每个标注shape for shape in label_data[shapes]: label_name shape[label] points shape[points] shape_type shape[shape_type] # 将多边形点转换为COCO需要的segmentation格式一维数组 segmentation [np.array(points).flatten().tolist()] # 计算包围框 [x_min, y_min, width, height] xs [p[0] for p in points] ys [p[1] for p in points] x_min, x_max min(xs), max(xs) y_min, y_max min(ys), max(ys) bbox_width x_max - x_min bbox_height y_max - y_min bbox [x_min, y_min, bbox_width, bbox_height] # 计算面积多边形面积 area poly_area(np.array(points)[:, 0], np.array(points)[:, 1]) coco_output[annotations].append({ id: annotation_id, image_id: image_id, category_id: category_name_to_id[label_name], segmentation: segmentation, area: float(area), bbox: bbox, iscrowd: 0 }) annotation_id 1 image_id 1 # 保存结果 os.makedirs(os.path.dirname(output_json_path), exist_okTrue) with open(output_json_path, w, encodingutf-8) as f: json.dump(coco_output, f, indent2) print(f转换完成保存至{output_json_path}) def poly_area(x, y): 计算多边形面积 return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1))) if __name__ __main__: # 示例用法 categories [cat, dog] # 请替换为你的实际类别 labelme_to_coco( labelme_dirdatas/train, output_json_pathdatas/annotations/instances_train.json, category_listcategories ) labelme_to_coco( labelme_dirdatas/val, output_json_pathdatas/annotations/instances_val.json, category_listcategories )运行这个脚本后你会在datas/annotations/文件夹下得到两个标准的COCO注解文件。同时建议把train和val文件夹下的所有图片统一移动到一个新的datas/images/文件夹下或者保持原样但确保JSON文件中的imagePath路径能被正确找到。COCO格式的JSON里记录的file_name最好是相对于数据集根目录的路径。这是最需要耐心和细心调试的一环转换成功后万里长征就走完一大半了。3. 注册与配置告诉模型“数据在哪里要学什么”数据准备好了现在得让Mask2Former框架知道我们这些数据的存在这就是“注册数据集”。同时我们还要通过配置文件告诉模型一些关键信息比如用哪个网络结构、学多少类、怎么学。这一步就像给厨师模型菜单数据和菜谱配置。注册数据集通常在启动训练的主文件里进行最常见的位置就是在train_net.py里。你不需要修改框架源码而是在调用训练代码前加上几行注册语句。具体来说在你自己的训练脚本或者直接修改train_net.py的开头部分在main函数之前添加如下代码from detectron2.data.datasets import register_coco_instances # 注册训练集 register_coco_instances( my_dataset_train, # 你给训练集起的名字后面配置里要用到 {}, # 元数据可以先留空字典或者像这样添加类别信息{thing_classes: [cat, dog]} datas/annotations/instances_train.json, # 注解文件路径 datas/train # 图片文件夹路径。如果图片在datas/images这里就写datas/images ) # 注册验证集 register_coco_instances( my_dataset_val, {}, datas/annotations/instances_val.json, datas/val )这里有个细节register_coco_instances的第四个参数是图片根目录。如果你的图片还在原来的train和val分开的文件夹里并且COCO JSON文件里的file_name包含了train/1.jpg这样的子路径那么第四个参数就应该只是datas。如果file_name直接是1.jpg那么第四个参数就需要是datas/train或datas/val。一定要匹配否则会找不到图片。我建议采用第一种方式即所有图片放在datas/imagesJSON里的file_name就是images/1.jpg注册时根目录写datas这样最清晰。注册完数据集接下来就是修改配置文件。Mask2Former的配置文件在configs/目录下结构清晰但文件众多。对于实例分割任务我们通常以configs/coco/instance-segmentation/下的配置文件为起点。比如我们复制一份maskformer2_R50_bs16_50ep.yaml并重命名为my_maskformer_config.yaml然后主要修改两个地方。首先找到DATASETS部分将其指向我们刚刚注册的数据集名字DATASETS: TRAIN: (my_dataset_train,) # 注意这里的括号和逗号表示这是一个元组 TEST: (my_dataset_val,) # 评估/验证时用的数据集其次找到MODEL部分下的MASK_FORMER或SEM_SEG_HEAD具体位置可能因版本略有不同修改分类头数量MODEL: MASK_FORMER: NUM_CLASSES: 3 # 修改为你的类别数11是背景类。例如你有猫和狗两类这里就填3。这里是个超级大坑Mask2Former的NUM_CLASSES到底应该设多少根据我的实测和经验在实例分割配置中这个值通常应该等于你的目标类别数量 1。这个“1”是给背景类留的。比如你只分割“猫”和“狗”两类物体那么NUM_CLASSES应该设为3。有些教程或代码里可能直接让你设成类别数如果不确定最好的方法是跑起来看报错信息或者去查看官方COCO配置的数值COCO有80类配置里设的是81。此外你还可以根据你的GPU显存调整批量大小SOLVER.IMS_PER_BATCH如果显存小比如8G可能需要从16调到4或2并相应地按比例增加训练轮数SOLVER.MAX_ITER以保证总“见到的”图片数差不多。学习率SOLVER.BASE_LR通常也需要随批量大小调整经验公式是新学习率 基础学习率 * (新批量大小 / 原批量大小)。4. 启动训练与调试按下启动键并解决路上的小麻烦万事俱备只欠一句训练命令。在终端中进入Mask2Former项目根目录执行类似下面的命令python train_net.py \ --config-file configs/coco/instance-segmentation/my_maskformer_config.yaml \ --num-gpus 1 \ OUTPUT_DIR ./output/my_first_training--num-gpus 1表示使用1张GPU如果你有多张可以改成相应数字。OUTPUT_DIR指定了模型权重和日志的输出目录建议每次都指定一个明确的、新的文件夹方便管理。按下回车后如果一切顺利你会看到终端开始刷屏显示加载数据集、构建模型、开始迭代训练的信息包括损失值下降、学习率变化等。这感觉非常棒但更常见的情况是你会遇到一些错误。别慌这是学习的必经之路。常见错误1KeyError: ‘XXX‘ is not in the field或找不到数据集。这几乎总是数据集注册或路径问题。请检查1) 注册数据集的名字是否和配置文件里写的完全一致大小写、下划线2) 注解文件路径和图片路径是否正确3) COCO JSON文件格式是否有效可以用python -m json.tool your_file.json验证一下JSON格式。常见错误2CUDA out of memory。显存炸了。首先尝试减小SOLVER.IMS_PER_BATCH。如果已经减小到1还是不行可以尝试减小输入图片的分辨率。在配置文件中搜索INPUT.MIN_SIZE_TRAIN和INPUT.MAX_SIZE_TRAIN将其值调小例如从[640, 672, 704, 736, 768, 800]调到[320, 352, 384, 416, 448, 480]。同时也可以使用梯度累积来模拟更大的批量但这需要修改训练代码对新手稍复杂。常见错误3损失为NaN或变得巨大。这可能是学习率太高了。将SOLVER.BASE_LR调低一个数量级试试例如从0.0001调到0.00001。也可能是数据标注有问题比如出现了坐标点超出图像范围、面积为0的多边形等。检查你的数据转换脚本确保生成的标注是合理的。训练开始后建议使用TensorBoard来监控训练过程。Detectron2默认会记录日志到输出目录下的metrics.json或events.out文件。你可以在另一个终端运行tensorboard --logdir ./output/my_first_training然后在浏览器打开提示的地址就能看到损失曲线、学习率、验证集精度等漂亮的可视化图表了。这不仅能帮你判断模型是否在正常学习还能为调整超参数提供依据。5. 模型推理与应用看看你的“孩子”学得怎么样训练了几个小时甚至几天模型终于保存了检查点通常保存在OUTPUT_DIR下如model_0004999.pth。是时候检验成果了Mask2Former项目通常自带一个演示脚本demo.py我们可以用它来对单张图片、多张图片甚至视频进行预测。一个基本的预测命令如下python demo/demo.py \ --config-file configs/coco/instance-segmentation/my_maskformer_config.yaml \ --input datas/val/100.jpg \ --output ./prediction_output \ --opts MODEL.WEIGHTS output/my_first_training/model_final.pth--input可以是一张图片的路径、一个包含图片的文件夹路径如datas/val/*.jpg、或者一个视频文件。--output指定一个目录预测结果带分割蒙版和框的图片会保存到这里。--opts MODEL.WEIGHTS这是最关键的一项指定你要使用的训练好的模型权重文件路径。通常最终模型叫model_final.pth中间检查点像model_0004999.pth也可以用。--confidence-threshold可以加上这个参数比如--confidence-threshold 0.5只显示置信度高于0.5的预测结果这样画面会更干净。运行命令后如果成功你会在输出目录看到一张和输入图同名的图片上面用不同颜色标出了模型识别和分割出的物体并附上了类别标签和置信度。第一次看到自己训练的模型准确地圈出图片中的目标时那种成就感是无与伦比的如果预测结果不理想比如什么都检测不出来或者分割得一塌糊涂别灰心。回过头检查训练数据量是否足够至少每类几百张标注质量高吗边界框和分割多边形是否精确训练轮数MAX_ITER够不够模型是否收敛看TensorBoard里的训练损失是否已经平稳在低位有时候用预训练权重进行微调会比完全从头训练效果好得多。你可以在训练命令中通过MODEL.WEIGHTS参数加载一个在COCO上预训练的模型例如官方提供的model_final_94dc52.pkl这样模型已经有了基本的视觉特征提取能力学起来会快很多效果也更好。最后把训练好的模型集成到你自己的应用里才是终极目标。你需要了解如何加载配置和模型并进行前向推理。这里提供一个最简单的Python脚本示例import cv2 import torch from detectron2.config import get_cfg from detectron2.engine import DefaultPredictor from detectron2.utils.visualizer import Visualizer from detectron2.data import MetadataCatalog # 1. 加载配置 cfg get_cfg() cfg.merge_from_file(configs/coco/instance-segmentation/my_maskformer_config.yaml) cfg.MODEL.WEIGHTS output/my_first_training/model_final.pth cfg.MODEL.DEVICE cuda if torch.cuda.is_available() else cpu # 设置预测阈值 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST 0.5 cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON False cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON True # 2. 创建预测器 predictor DefaultPredictor(cfg) # 3. 读取图片并预测 image cv2.imread(your_test_image.jpg) outputs predictor(image) # 4. 可视化结果 # 获取数据集的元数据主要是类别名 metadata MetadataCatalog.get(cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else __unused) v Visualizer(image[:, :, ::-1], metadata, scale1.2) # BGR转RGB out v.draw_instance_predictions(outputs[instances].to(cpu)) result_image out.get_image()[:, :, ::-1] # RGB转回BGR # 5. 保存或显示 cv2.imwrite(prediction_result.jpg, result_image)这个过程从无到有虽然步骤繁多但每一步都有其意义。环境配置是打好地基数据准备是准备食材注册配置是制定菜谱训练是烹饪过程而推理和应用就是品尝和上菜。我刚开始的时候在数据转换和配置修改上卡了最久往往就是一个路径错误或者一个参数理解偏差。希望我分享的这些具体步骤和踩过的坑能帮你更平滑地完成这段旅程。记住遇到报错多查查GitHub的issue很多问题别人都遇到过。当你第一次用自己的数据跑通整个流程看到模型准确分割出目标时那种快乐会让你觉得所有的折腾都是值得的。