从零开始：使用mmsegmentation训练自定义数据集的全流程指南

1. 环境准备与安装指南第一次接触mmsegmentation时最头疼的就是环境配置。记得我刚开始用的时候光是解决CUDA和PyTorch版本兼容问题就折腾了一整天。现在把踩过的坑都总结出来让你10分钟搞定环境搭建。核心依赖清单Python 3.7/3.8实测3.9会有兼容性问题PyTorch 1.6CUDA 10.1/10.2/11.1必须与显卡驱动匹配mmcv-full注意要选对应PyTorch和CUDA的版本具体操作步骤# 创建虚拟环境建议用conda管理 conda create -n mmseg python3.8 -y conda activate mmseg # 安装PyTorch以CUDA 10.1为例 pip install torch1.8.1cu101 torchvision0.9.1cu101 -f https://download.pytorch.org/whl/torch_stable.html # 安装mmcv-full版本查询表见mmcv官方文档 pip install mmcv-full1.4.0 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.8/index.html # 安装mmsegmentation git clone https://github.com/open-mmlab/mmsegmentation.git cd mmsegmentation pip install -e .常见问题解决方案如果遇到CUDA out of memory错误尝试减小batch_size修改config中的samples_per_gpu使用SyncBN替换BN在模型config中修改安装mmcv-full时报错检查PyTorch和CUDA版本是否匹配尝试从源码编译MMCV_WITH_OPS1 pip install mmcv-full提示建议使用Docker镜像规避环境问题官方提供了预配置好的镜像docker pull openmmlab/mmsegmentation:latest2. 数据标注与格式转换实战用labelme标注数据时我发现几个新手容易忽略的关键点标注时要闭合多边形最后一点必须与第一点重合同类物体多个实例需要分开标注建议用英文命名类别避免编码问题标注完成后典型的文件结构是这样的my_dataset/ ├── img_dir/ │ ├── train/ │ │ ├── 001.jpg │ │ └── 002.jpg │ └── val/ │ └── 003.jpg └── ann_dir/ ├── train/ │ ├── 001.png │ └── 002.png └── val/ └── 003.png转换脚本的核心逻辑解析def labelme2mask(json_path, save_dir): # 读取JSON标注文件 with open(json_path) as f: data json.load(f) # 创建空白mask画布 height, width data[imageHeight], data[imageWidth] mask np.zeros((height, width), dtypenp.uint8) # 遍历所有标注区域 for shape in data[shapes]: label shape[label] points np.array(shape[points], dtypenp.int32) # 将标注区域填充为类别ID cv2.fillPoly(mask, [points], colorclass_dict[label]) # 保存为PNG格式 cv2.imwrite(os.path.join(save_dir, os.path.basename(json_path).replace(.json,.png)), mask)实际项目中的经验技巧小目标处理对小于50x50像素的物体建议标注时适当扩大边界类别平衡如果某些类别样本过少可以在dataset配置中设置repeat_times边缘模糊问题在train_pipeline中添加RandomBlur增强3. 配置文件深度定制指南mmsegmentation的配置文件系统非常灵活但也容易让人困惑。我拆解一个典型的config文件结构configs/ └── my_model/ ├── deeplabv3plus_r50-d8.py # 模型结构 ├── my_dataset.py # 数据集配置 ├── schedule_20k.py # 训练计划 └── default_runtime.py # 运行时配置关键参数调整策略参数推荐值作用lr0.01 (bs8)基础学习率需随batch_size调整crop_size(512,512)裁剪尺寸越大显存消耗越高samples_per_gpu2-8根据GPU显存调整workers_per_gpu4-8数据加载线程数自定义数据集配置示例# my_dataset.py dataset_type MyDataset data_root data/my_dataset train_pipeline [ dict(typeLoadImageFromFile), dict(typeLoadAnnotations), dict(typeRandomFlip, prob0.5), dict(typeNormalize, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375]), dict(typeDefaultFormatBundle), dict(typeCollect, keys[img, gt_semantic_seg]) ] data dict( traindict( typedataset_type, data_rootdata_root, img_dirimg_dir/train, ann_dirann_dir/train, pipelinetrain_pipeline), valdict(...), testdict(...) )4. 训练与调优全流程启动训练的命令其实很简单python tools/train.py configs/my_model/deeplabv3plus_r50-d8.py --work-dir work_dir但真正影响效果的是这些细节学习率策略在schedule_20k.py中调整optimizer dict(typeSGD, lr0.01, momentum0.9, weight_decay0.0005) lr_config dict(policypoly, power0.9, min_lr1e-4, by_epochFalse)模型保存修改default_runtime.pycheckpoint_config dict(interval2000) # 每2000次迭代保存一次 evaluation dict(interval2000, metricmIoU) # 评估指标数据增强推荐组合方案train_pipeline [ ... dict(typeRandomRotate, prob0.5, degree(-30, 30)), dict(typeColorJitter, brightness0.5, contrast0.5, saturation0.5), dict(typeGaussianBlur, sigma_range(0.1, 2.0), prob0.5) ]训练过程监控技巧使用TensorBoard查看实时指标tensorboard --logdir work_dir --port 6006中断后恢复训练python tools/train.py configs/my_model/deeplabv3plus_r50-d8.py \ --resume-from work_dir/latest.pth5. 模型部署与效果验证训练完成后用这个脚本测试单张图片from mmseg.apis import inference_segmentor, init_segmentor config_file configs/my_model/deeplabv3plus_r50-d8.py checkpoint_file work_dir/latest.pth model init_segmentor(config_file, checkpoint_file, devicecuda:0) result inference_segmentor(model, test.jpg) model.show_result(test.jpg, result, out_fileresult.jpg)实际项目中遇到的性能优化点显存不足尝试这些方案启用SyncBN多GPU训练时减小crop_size如从512x512降到256x256使用FP16训练需PyTorch 1.6预测速度慢优化方向# 修改config中的test_pipeline test_pipeline [ dict(typeLoadImageFromFile), dict(typeMultiScaleFlipAug, img_scale(1024, 1024), # 减小尺寸 flipFalse, # 关闭测试时翻转 transforms[...]) ]边缘锯齿问题后处理方案import cv2 result inference_segmentor(model, img) result cv2.medianBlur(result, ksize3) # 中值滤波最后分享一个实用技巧把常用配置封装成函数比如我的quick_test.py包含def visualize(model, img_path): result inference_segmentor(model, img_path) return model.show_result(img_path, result, opacity0.5) def batch_predict(model, img_dir): # 批量预测并保存结果 ...