MMSegmentation 模型训练与测试全指南

MMSegmentation 模型训练与测试全指南

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

前言

MMSegmentation 是一个基于 PyTorch 的语义分割工具库，提供了丰富的模型实现和训练测试工具。本文将详细介绍如何使用 MMSegmentation 进行模型的训练和测试，涵盖单机单卡、多卡并行以及集群环境下的操作流程。

单GPU训练与测试

单GPU训练

在单GPU环境下训练模型是最基础的操作方式，适用于小规模数据集或模型调试阶段。

基本命令格式：

python tools/train.py ${配置文件} [可选参数]

关键参数说明：

• --work-dir: 指定工作目录路径，用于保存训练日志和模型检查点
• --amp: 启用自动混合精度训练，可减少显存占用并加速训练
• --resume: 从最近的检查点恢复训练
• --cfg-options: 动态修改配置参数，格式为键值对

CPU训练模式： 当需要在CPU上训练时，只需设置环境变量：

export CUDA_VISIBLE_DEVICES=-1

单GPU测试

模型训练完成后，可以使用测试脚本评估模型性能。

基本命令格式：

python tools/test.py ${配置文件} ${模型权重} [可选参数]

可视化参数：

• --show: 实时显示预测结果
• --show-dir: 保存可视化结果到指定目录
• --wait-time: 可视化间隔时间(秒)

CPU测试模式： 与训练类似，通过设置环境变量切换到CPU模式：

export CUDA_VISIBLE_DEVICES=-1

多GPU与分布式训练测试

多GPU训练

对于大规模数据集或复杂模型，多GPU并行训练可以显著缩短训练时间。

基本命令格式：

sh tools/dist_train.sh ${配置文件} ${GPU数量} [可选参数]

注意事项：

1. 默认使用所有可见GPU
2. 日志和检查点会自动保存在配置文件名对应的目录下
3. 建议使用符号链接管理工作目录

多GPU测试

基本命令格式：

sh tools/dist_test.sh ${配置文件} ${检查点} ${GPU数量} [可选参数]

单机多任务管理

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

# 任务1
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 sh tools/dist_train.sh ${配置} 4
# 任务2 
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 sh tools/dist_train.sh ${配置} 4

多机分布式训练

在多台机器上分布式训练需要指定主节点地址和端口：

主节点：

NNODES=2 NODE_RANK=0 PORT=29500 MASTER_ADDR=主节点IP sh tools/dist_train.sh ${配置} ${GPUS}

从节点：

NNODES=2 NODE_RANK=1 PORT=29500 MASTER_ADDR=主节点IP sh tools/dist_train.sh ${配置} ${GPUS}

Slurm集群环境管理

Slurm训练任务

在Slurm管理的集群上，可以使用专用脚本提交训练任务：

基本命令格式：

GPUS=${数量} sh tools/slurm_train.sh ${分区} ${任务名} ${配置} [参数]

Slurm测试任务

基本命令格式：

GPUS=${数量} sh tools/slurm_test.sh ${分区} ${任务名} ${配置} ${检查点} [参数]

端口冲突解决方案

在Slurm环境下运行多个任务时，可通过以下方式避免端口冲突：

1. 通过--cfg-options动态设置端口：

--cfg-options env_cfg.dist_cfg.port=29500

2. 直接修改配置文件：

env_cfg = dict(dist_cfg=dict(backend='nccl', port=29500))

3. 通过环境变量设置：

MASTER_PORT=29500 sh tools/slurm_train.sh ...

结果保存与可视化

基础结果保存

测试时可通过--out参数保存预测结果：

python tools/test.py ${配置} ${检查点} --out ${输出目录}

无标注数据测试

当测试数据没有标注时，需要修改配置：

test_evaluator = dict(
    type='IoUMetric',
    format_only=True,
    output_dir='输出路径')
    
test_dataloader = dict(
    dataset=dict(
        pipeline=[
            dict(type='LoadImageFromFile'),
            # 移除LoadAnnotations
            dict(type='Resize', scale=(2048, 512), keep_ratio=True),
            dict(type='PackSegInputs')
        ]))

Cityscapes数据集特殊处理

对于Cityscapes数据集，推荐使用专用评估器：

test_evaluator = dict(
    type='CityscapesMetric',
    format_only=True,
    keep_results=True,
    output_dir='输出路径')

总结

本文全面介绍了MMSegmentation在不同环境下的训练和测试方法，从单机单卡到分布式集群环境，涵盖了各种实际应用场景。掌握这些方法可以帮助研究人员和开发者高效地完成语义分割模型的训练和评估工作。

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