Mask2Former实战：从零部署到自定义数据集训练全攻略

1. 环境配置从零搭建Mask2Former开发环境第一次接触Mask2Former时我也被复杂的依赖关系搞得头大。经过多次尝试发现用官方推荐的Docker镜像确实能省去不少麻烦。这里分享我的实测经验使用NVIDIA官方提供的PyTorch 1.10镜像作为基础再叠加Detectron2和Mask2Former的依赖成功率最高。具体操作步骤# 拉取基础镜像 docker pull nvcr.io/nvidia/pytorch:21.08-py3 # 启动容器时记得挂载数据集目录 docker run -it --gpus all -v /本地数据集路径:/container/dataset nvcr.io/nvidia/pytorch:21.08-py3进入容器后按这个顺序安装关键组件# 先装Detectron2 git clone https://github.com/facebookresearch/detectron2.git cd detectron2 pip install -e . # 再装PanopticAPI等辅助工具 pip install githttps://github.com/cocodataset/panopticapi.git pip install githttps://github.com/mcordts/cityscapesScripts.git # 最后安装Mask2Former cd .. git clone https://github.com/facebookresearch/Mask2Former.git cd Mask2Former pip install -r requirements.txt有个坑要特别注意Pixel Decoder的CUDA扩展需要单独编译。我在三台不同配置的机器上都遇到过编译失败的情况后来发现是GCC版本不匹配。解决方法是用conda安装gxx_linux-64conda install gxx_linux-649.3.0 cd mask2former/modeling/pixel_decoder/ops python setup.py build install2. Demo测试验证环境是否正常工作装好环境不跑个Demo总觉得心里不踏实。Mask2Former官方提供了预训练模型我们可以用Cityscapes数据集快速验证。这里有个小技巧下载模型时用wget的--continue参数避免网络波动导致重头下载。# 下载预训练模型 mkdir -p weights wget --continue https://dl.fbaipublicfiles.com/maskformer/mask2former/coco/panoptic/maskformer2_swin_large_IN21k_384_bs16_100ep/model_final_f07440.pkl -P weights/ # 运行Demo python demo/demo.py --config-file configs/coco/panoptic-segmentation/swin/maskformer2_swin_large_IN21k_384_bs16_100ep.yaml \ --input input.jpg --output output.png \ --opts MODEL.WEIGHTS weights/model_final_f07440.pkl如果看到终端输出推理时间并且生成了带分割效果的output.png说明环境配置成功。我遇到过输出图片全黑的情况后来发现是OpenCV版本问题降级到4.5.5就正常了。3. 准备自定义数据集以ADE20K格式为例真实项目中我们往往需要处理自定义数据。Mask2Former支持多种数据集格式实测下来ADE20K格式最容易上手。目录结构应该是这样的my_dataset/ ├── images/ │ ├── train/ │ │ ├── 0001.jpg │ │ └── ... │ └── val/ │ ├── 0001.jpg │ └── ... └── annotations/ ├── train/ │ ├── 0001.png │ └── ... └── val/ ├── 0001.png └── ...标注图像需要特别注意颜色编码背景必须用[0,0,0]表示每个类别用[R,G,B]表示且RGB如[1,1,1],[2,2,2]建议用LabelMe标注后用这个Python脚本转换import numpy as np from PIL import Image def convert_labelme_to_ade20k(labelme_mask_path, output_path): mask np.array(Image.open(labelme_mask_path)) # 将标注值转换为单通道 ade20k_mask mask[:,:,0].astype(np.uint8) Image.fromarray(ade20k_mask).save(output_path)对于二分类任务只有背景和前景我发现直接使用原始标注比用prepare_ade20k_sem_seg.py更可靠。因为转换脚本会把标签值减1导致背景变成255可能引发训练异常。4. 模型训练关键配置详解准备好数据集后需要修改配置文件。以Swin-Transformer为例主要调整这几个关键参数数据集注册在train_net.py前添加from detectron2.data.datasets import register_coco_panoptic_separated register_coco_panoptic_separated( my_dataset_train, metadata{}, image_rootpath/to/images/train, panoptic_rootpath/to/annotations_detectron2/train, panoptic_json, # 语义分割留空 gt_json, # 语义分割留空 instance_root, # 语义分割留空 )修改configs/swin/maskformer2_swin_base_384_bs16_50ep.yamlMODEL: SEM_SEG_HEAD: NUM_CLASSES: 2 # 背景前景 WEIGHTS: pretrained/swin_base_patch4_window12_384.pth SOLVER: MAX_ITER: 10000 # 根据数据集大小调整 CHECKPOINT_PERIOD: 1000 IMS_PER_BATCH: 2 # 根据GPU显存调整 TEST: EVAL_PERIOD: 1000 OUTPUT_DIR: output/my_dataset_swinB_bs2启动训练时建议用分布式训练即使只有一张卡python train_net.py \ --config-file configs/swin/maskformer2_swin_base_384_bs16_50ep.yaml \ --num-gpus 1 \ OUTPUT_DIR output/my_dataset_swinB_bs2训练过程中常见的几个问题Loss不下降检查学习率默认0.0001可能太大尝试减小到1e-5GPU内存不足减小IMS_PER_BATCH或缩小图片尺寸验证指标波动大增加EVAL_PERIOD确保每次验证有足够样本5. 模型推理与部署训练完成后用这个脚本进行单张图片推理from detectron2.config import get_cfg from mask2former import add_maskformer2_config cfg get_cfg() add_maskformer2_config(cfg) cfg.merge_from_file(configs/swin/maskformer2_swin_base_384_bs16_50ep.yaml) cfg.MODEL.WEIGHTS output/my_dataset_swinB_bs2/model_final.pth predictor DefaultPredictor(cfg) outputs predictor(im)对于生产环境部署我推荐将模型导出为TorchScript格式from mask2former.export import export_model export_model( cfg, output/my_dataset_swinB_bs2/model_final.pth, deploy/model.pt, size_divisibility32 )实际项目中我发现用Triton Inference Server部署性能最好。编写config.pbtxt时要注意platform: pytorch_libtorch max_batch_size: 4 input [ { name: images data_type: TYPE_FP32 dims: [3, -1, -1] } ] output [ { name: sem_seg data_type: TYPE_FP32 dims: [2, -1, -1] # 与NUM_CLASSES对应 } ]6. 性能优化技巧经过多个项目实践我总结出这些提升Mask2Former效率的方法混合精度训练在配置文件中添加TRAINER: default TEST: PRECISE_BN: ENABLED: True SOLVER: AMP: ENABLED: True数据加载优化使用更快的图片解码库pip install pyvips # 然后在数据加载器中指定 DATALOADER: IMAGE_BACKEND: pyvips模型裁剪对于二分类任务可以减小Transformer的隐藏层维度MODEL: MASK_FORMER: HIDDEN_DIM: 256 # 默认512 NUM_OBJECT_QUERIES: 50 # 默认100缓存数据集首次训练时添加python train_net.py --use-dataset-cache这些优化能让训练速度提升2-3倍特别是在小数据集上效果明显。我在一个5000张图片的项目中训练时间从12小时缩短到了4小时。