MMSegmentation 模型训练与测试全指南

MMSegmentation 模型训练与测试全指南

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation

前言

MMSegmentation 作为开源的语义分割工具库，提供了完整的模型训练和测试流程。本文将详细介绍如何使用 MMSegmentation 进行单机单卡、多卡以及集群环境下的模型训练与测试，帮助开发者快速上手并高效利用计算资源。

单GPU训练与测试

单卡训练

使用 tools/train.py 脚本可以轻松启动单卡训练任务：

python tools/train.py configs/pspnet/pspnet_r50-d8_4xb4-80k_ade20k-512x512.py

常用参数说明

• --work-dir: 指定工作目录，覆盖配置文件中的默认设置
• --amp: 启用自动混合精度训练，可减少显存占用
• --resume: 自动从工作目录中恢复最新检查点
• --cfg-options: 动态修改配置参数，格式为 key=value

重要概念区分

1. --resume 与 load_from 的区别：

   • --resume 用于意外中断后恢复训练
   • load_from 用于指定预训练模型进行微调

2. 在CPU上训练：

export CUDA_VISIBLE_DEVICES=-1
python tools/train.py ...

单卡测试

测试脚本 tools/test.py 的基本用法：

python tools/test.py configs/pspnet/pspnet_r50-d8_4xb4-80k_ade20k-512x512.py checkpoint.pth

测试功能选项

• --show: 实时显示预测结果
• --show-dir: 保存可视化分割结果
• --wait-time: 显示间隔时间(秒)
• --tta: 启用测试时数据增强

多GPU与分布式训练

多卡训练

使用分布式训练脚本 tools/dist_train.sh：

sh tools/dist_train.sh configs/pspnet/pspnet_r50-d8_4xb4-80k_ade20k-512x512.py 8

工作目录建议

MMSegmentation 会自动按照配置文件结构组织工作目录。如需自定义位置，建议使用符号链接：

ln -s /your/custom/path ./work_dirs

多卡测试

分布式测试脚本 tools/dist_test.sh 用法：

sh tools/dist_test.sh config_file checkpoint.pth 8

单机多任务管理

当在同一台机器上运行多个训练任务时，需要指定不同端口避免冲突：

CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 sh tools/dist_train.sh config1.py 4
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 sh tools/dist_train.sh config2.py 4

Slurm集群管理

Slurm训练任务

使用 slurm_train.sh 脚本提交训练任务：

GPUS=4 sh tools/slurm_train.sh dev pspnet configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py

Slurm测试任务

类似地，使用 slurm_test.sh 提交测试任务：

GPUS=4 sh tools/slurm_test.sh dev pspnet_test config_file checkpoint.pth

端口设置方法

在Slurm环境下有三种设置端口的方式：

1. 通过 --cfg-options 动态设置（推荐）
2. 直接修改配置文件
3. 使用环境变量 MASTER_PORT

结果保存与提交

基础保存功能

测试时保存预测结果：

python tools/test.py config_file checkpoint.pth --out results_dir

或通过配置文件设置：

test_evaluator = dict(type='IoUMetric', output_dir='results_dir')

无标注数据测试

对于无标注数据，设置 format_only=True 并移除标注加载：

test_evaluator = dict(
    type='IoUMetric',
    format_only=True,
    output_dir='results_dir')

Cityscapes数据集提交

使用 CityscapesMetric 生成符合官方要求的提交文件：

test_evaluator = dict(
    type='CityscapesMetric',
    format_only=True,
    output_dir='cityscapes_results')

最佳实践建议

1. 对于大型数据集，优先使用分布式训练
2. 测试时合理利用 --tta 选项提升模型性能
3. 使用符号链接管理工作目录，保持项目结构清晰
4. 在集群环境中，提前规划好端口分配方案
5. 提交竞赛结果前，仔细检查输出格式要求

通过本文介绍的各种训练和测试方法，开发者可以根据自身硬件条件和项目需求，灵活选择最适合的方案，充分发挥 MMSegmentation 的强大功能。

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation