mmsegmentation教程v0.20.2

1,595 阅读7分钟

教程1:训练与测试模型

训练一个模型

MMSegmentation 可以执行分布式训练和非分布式训练,分别用 MMDistributedDataParallel 和 MMDataParallel 命令。 所有的输出(日志 log 和检查点 checkpoints )将被保存到工作路径文件夹里,它可以通过配置文件里的 work_dir 指定。 在一定迭代轮次后,我们默认在验证集上评估模型表现。您可以在训练配置文件中添加间隔参数来改变评估间隔。

evaluation = dict(interval=4000)  # 每4000 iterations 评估一次模型的性能

Important : 在配置文件里的默认学习率是针对4卡 GPU 和2张图/GPU (此时 batchsize = 4x2 = 8)来设置的。 同样,您也可以使用8卡 GPU 和 1张图/GPU 的设置,因为所有的模型均使用 cross-GPU 的 SyncBN 模式。\

我们可以在训练速度和 GPU 显存之间做平衡。当模型或者 Batch Size 比较大的时,可以传递--cfg-options model.backbone.with_cp=True ,使用 with_cp 来节省显存,但是速度会更慢,因为原先使用 ith_cp 时,是逐层反向传播(Back Propagation, BP),不会保存所有的梯度。

使用单卡 GPU 训练

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

如果您想在命令里定义工作文件夹路径,您可以添加一个参数--work-dir ${YOUR_WORK_DIR}

使用多卡 GPU 训练

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

可选参数可以为:

  • --no-validate (不推荐): 训练时代码库默认会在每 k 轮迭代后在验证集上进行评估,如果不需评估使用命令 --no-validate
  • --work-dir ${工作路径}: 在配置文件里重写工作路径文件夹
  • --resume-from ${检查点文件}: 继续使用先前的检查点 (checkpoint) 文件(可以继续训练过程)
  • --load-from ${检查点文件}: 从一个检查点 (checkpoint) 文件里加载权重(对另一个任务进行精调)

resume-from 和 load-from 的区别:

  • resume-from 加载出模型权重和优化器状态包括迭代轮数等
  • load-from 仅加载模型权重,从第0轮开始训练

使用多个机器训练(暂时不掌握)

如果您在一个集群上以slurm 运行 MMSegmentation, 您可以使用脚本 slurm_train.sh(这个脚本同样支持单个机器的训练)。

[GPUS=${GPU 数量}] ./tools/slurm_train.sh ${分区} ${任务名称} ${配置文件} --work-dir ${工作路径}

这里是在 dev 分区里使用16块 GPU 训练 PSPNet 的例子。

GPUS=16 ./tools/slurm_train.sh dev pspr50 configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py /nfs/xxxx/psp_r50_512x1024_40ki_cityscapes

您可以查看 slurm_train.sh 以熟悉全部的参数与环境变量。

如果您多个机器已经有以太网连接, 您可以参考 PyTorch launch utility 。 若您没有像 InfiniBand 这样高速的网络连接,多机器训练通常会比较慢。

在单个机器上启动多个任务

如果您在单个机器上启动多个任务,例如在8卡 GPU 的一个机器上有2个4卡 GPU 的训练任务,您需要特别对每个任务指定不同的端口(默认为29500)来避免通讯冲突。 否则,将会有报错信息 RuntimeError: Address already in use。 如果您使用命令 dist_train.sh 来启动一个训练任务,您可以在命令行的用环境变量 PORT 设置端口。

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

如果您使用命令 slurm_train.sh 来启动训练任务,您可以在命令行的用环境变量 MASTER_PORT 设置端口。

MASTER_PORT=29500 ./tools/slurm_train.sh ${分区} ${任务名称} ${配置文件}
MASTER_PORT=29501 ./tools/slurm_train.sh ${分区} ${任务名称} ${配置文件}

使用预训练模型推理

我们提供测试脚本来评估完整数据集(Cityscapes, PASCAL VOC, ADE20k 等)上的结果,同时为了使其他项目的整合更容易,也提供一些高级 API。

测试一个数据集

  • 单卡 GPU
  • 单节点多卡 GPU
  • 多节点

您可以使用以下命令来测试一个数据集。

# 单卡 GPU 测试
python tools/test.py ${配置文件} ${检查点文件} [--out ${结果文件:RESULT_FILE}] [--eval ${评估指标:EVAL_METRICS}] [--show]

# 多卡GPU 测试
./tools/dist_test.sh ${配置文件} ${检查点文件} ${GPU数目} [--out ${结果文件:RESULT_FILE}] [--eval ${评估指标:EVAL_METRICS}]

可选参数:

  • RESULT_FILE: pickle 格式的输出结果的文件名,如果不专门指定,结果将不会被专门保存成文件。(MMseg v0.17 之后,args.out 将只会保存评估时的中间结果或者是分割图的保存路径。)
  • EVAL_METRICS: 在结果里将被评估的指标。这主要取决于数据集, mIoU 对于所有数据集都可获得,像 Cityscapes 数据集可以通过 cityscapes 命令来专门评估,就像标准的 mIoU一样。
  • --show: 如果被指定,分割结果将会在一张图像里画出来并且在另一个窗口展示。它仅仅是用来调试与可视化,并且仅针对单卡 GPU 测试。请确认 GUI 在您的环境里可用,否则您也许会遇到报错 cannot connect to X server
  • --show-dir: 如果被指定,分割结果将会在一张图像里画出来并且保存在指定文件夹里。它仅仅是用来调试与可视化,并且仅针对单卡GPU测试。使用该参数时,您的环境不需要 GUI。
  • --eval-options: 评估时的可选参数,当设置 efficient_test=True 时,它将会保存中间结果至本地文件里以节约 CPU 内存。请确认您本地硬盘有足够的存储空间(大于20GB)。(MMseg v0.17 之后,efficient_test 不再生效,我们重构了 test api,通过使用一种渐近式的方式来提升评估和保存结果的效率。)

例子:

假设您已经下载检查点文件至文件夹 checkpoints/ 里。

  1. 测试 PSPNet 并可视化结果。按下任何键会进行到下一张图
python tools/test.py configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    --show
  1. 测试 PSPNet 并保存画出的图以便于之后的可视化
python tools/test.py configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    --show-dir psp_r50_512x1024_40ki_cityscapes_results
  1. 在数据集 PASCAL VOC (不保存测试结果) 上测试 PSPNet 并评估 mIoU
python tools/test.py configs/pspnet/pspnet_r50-d8_512x1024_20k_voc12aug.py \
    checkpoints/pspnet_r50-d8_512x1024_20k_voc12aug_20200605_003338-c57ef100.pth \
    --eval mAP
  1. 使用4卡 GPU 测试 PSPNet,并且在标准 mIoU 和 cityscapes 指标里评估模型
./tools/dist_test.sh configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    4 --out results.pkl --eval mIoU cityscapes

注意:在 cityscapes mIoU 和我们的 mIoU 指标会有一些差异 (~0.1%) 。因为 cityscapes 默认是根据类别样本数的多少进行加权平均,而我们对所有的数据集都是采取直接平均的方法来得到 mIoU。

  1. 在 cityscapes 数据集上4卡 GPU 测试 PSPNet, 并生成 png 文件以便提交给官方评估服务器
    首先,在配置文件里添加内容: configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py
data = dict(
    test=dict(
        img_dir='leftImg8bit/test',
        ann_dir='gtFine/test'))

随后,进行测试。

./tools/dist_test.sh configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    4 --format-only --eval-options "imgfile_prefix=./pspnet_test_results"

您会在文件夹 ./pspnet_test_results 里得到生成的 png 文件。 您也许可以运行 zip -r results.zip pspnet_test_results/ 并提交 zip 文件给 evaluation server 。 6. 在 Cityscapes 数据集上使用 CPU 高效内存选项来测试 DeeplabV3+ mIoU 指标 (没有保存测试结果)

python tools/test.py \
configs/deeplabv3plus/deeplabv3plus_r18-d8_512x1024_80k_cityscapes.py \
deeplabv3plus_r18-d8_512x1024_80k_cityscapes_20201226_080942-cff257fe.pth \
--eval-options efficient_test=True \
--eval mIoU

使用 pmap 可查看 CPU 内存情况, efficient_test=True 会使用约 2.25GB 的 CPU 内存, efficient_test=False 会使用约 11.06GB 的 CPU 内存。 这个可选参数可以节约很多 CPU 内存。(MMseg v0.17 之后, efficient_test 参数将不再生效, 我们使用了一种渐近的方式来更加有效快速地评估和保存结果。)

  1. 在 LoveDA 数据集上1卡 GPU 测试 PSPNet, 并生成 png 文件以便提交给官方评估服务器 首先,在配置文件里添加内容: configs/pspnet/pspnet_r50-d8_512x512_80k_loveda.py,
data = dict(
    test=dict(
        img_dir='img_dir/test',
        ann_dir='ann_dir/test'))

随后,进行测试。

python ./tools/test.py configs/pspnet/pspnet_r50-d8_512x512_80k_loveda.py \
    checkpoints/pspnet_r50-d8_512x512_80k_loveda_20211104_155728-88610f9f.pth \
    --format-only --eval-options "imgfile_prefix=./pspnet_test_results"

您会在文件夹 ./pspnet_test_results 里得到生成的 png 文件。 您也许可以运行 zip -r -j Results.zip pspnet_test_results/ 并提交 zip 文件给 evaluation server 。

教程 2: 自定义数据集

通过重新组织数据来定制数据集

最简单的方法是将您的数据集进行转化,并组织成文件夹的形式。 如下的文件结构就是一个例子。

├── data
│   ├── my_dataset
│   │   ├── img_dir
│   │   │   ├── train
│   │   │   │   ├── xxx{img_suffix}
│   │   │   │   ├── yyy{img_suffix}
│   │   │   │   ├── zzz{img_suffix}
│   │   │   ├── val
│   │   ├── ann_dir
│   │   │   ├── train
│   │   │   │   ├── xxx{seg_map_suffix}
│   │   │   │   ├── yyy{seg_map_suffix}
│   │   │   │   ├── zzz{seg_map_suffix}
│   │   │   ├── val

一个训练对将由 img_dir/ann_dir 里同样首缀的文件组成。

如果给定 split 参数,只有部分在 img_dir/ann_dir 里的文件会被加载。 我们可以对被包括在 split 文本里的文件指定前缀。

除此以外,一个 split 文本如下所示:

xxx
zzz

只有

data/my_dataset/img_dir/train/xxx{img_suffix}data/my_dataset/img_dir/train/zzz{img_suffix}data/my_dataset/ann_dir/train/xxx{seg_map_suffix}data/my_dataset/ann_dir/train/zzz{seg_map_suffix} 将被加载。 (若split文本特殊,可通过修改custom文件进行加载)

注意:标注是跟图像同样的形状 (H, W),其中的像素值的范围是 [0, num_classes - 1]。 您也可以使用 pillow 的 'P' 模式去创建包含颜色的标注。

通过混合数据去定制数据集

MMSegmentation 同样支持混合数据集去训练。 当前它支持拼接 (concat) 和 重复 (repeat) 数据集。

重复数据集

我们使用 RepeatDataset 作为包装 (wrapper) 去重复数据集。 例如,假设原始数据集是 Dataset_A,为了重复它,配置文件如下:

dataset_A_train = dict(
        type='RepeatDataset',
        times=N,
        dataset=dict(  # 这是 Dataset_A 数据集的原始配置
            type='Dataset_A',
            ...
            pipeline=train_pipeline
        )
    )

拼接数据集

有2种方式去拼接数据集。

  1. 如果您想拼接的数据集是同样的类型,但有不同的标注文件, 您可以按如下操作去拼接数据集的配置文件:

    1. 您也许可以拼接两个标注文件夹 ann_dir

      dataset_A_train = dict(
          type='Dataset_A',
          img_dir = 'img_dir',
          ann_dir = ['anno_dir_1', 'anno_dir_2'],
          pipeline=train_pipeline
      )
      
    2. 您也可以去拼接两个 split 文件列表

      dataset_A_train = dict(
          type='Dataset_A',
          img_dir = 'img_dir',
          ann_dir = 'anno_dir',
          split = ['split_1.txt', 'split_2.txt'],
          pipeline=train_pipeline
      )
      
    3. 您也可以同时拼接 ann_dir 文件夹和 split 文件列表

      dataset_A_train = dict(
          type='Dataset_A',
          img_dir = 'img_dir',
          ann_dir = ['anno_dir_1', 'anno_dir_2'],
          split = ['split_1.txt', 'split_2.txt'],
          pipeline=train_pipeline
      )
      

      在这样的情况下, ann_dir_1 和 ann_dir_2 分别对应于 split_1.txt 和 split_2.txt

  2. 如果您想拼接不同的数据集,您可以如下去拼接数据集的配置文件:

    dataset_A_train = dict()
    dataset_B_train = dict()
    
    data = dict(
        imgs_per_gpu=2,
        workers_per_gpu=2,
        train = [
            dataset_A_train,
            dataset_B_train
        ],
        val = dataset_A_val,
        test = dataset_A_test
        )
    

一个更复杂的例子如下:分别重复 Dataset_A 和 Dataset_B N 次和 M 次,然后再去拼接重复后的数据集

dataset_A_train = dict(
    type='RepeatDataset',
    times=N,
    dataset=dict(
        type='Dataset_A',
        ...
        pipeline=train_pipeline
    )
)
dataset_A_val = dict(
    ...
    pipeline=test_pipeline
)
dataset_A_test = dict(
    ...
    pipeline=test_pipeline
)
dataset_B_train = dict(
    type='RepeatDataset',
    times=M,
    dataset=dict(
        type='Dataset_B',
        ...
        pipeline=train_pipeline
    )
)
data = dict(
    imgs_per_gpu=2,
    workers_per_gpu=2,
    train = [
        dataset_A_train,
        dataset_B_train
    ],
    val = dataset_A_val,
    test = dataset_A_test
)

教程 3: 自定义数据流程

数据流程的设计

按照通常的惯例,我们使用 Dataset 和 DataLoader 做多线程的数据加载。Dataset 返回一个数据内容的字典,里面对应于模型前传方法的各个参数。 因为在语义分割中,输入的图像数据具有不同的大小,我们在 MMCV 里引入一个新的 DataContainer 类别去帮助收集和分发不同大小的输入数据。

更多细节,请查看这里 。

数据的准备流程和数据集是解耦的。通常一个数据集定义了如何处理标注数据(annotations)信息,而一个数据流程定义了准备一个数据字典的所有步骤。一个流程包括了一系列操作,每个操作里都把一个字典作为输入,然后再输出一个新的字典给下一个变换操作。

这些操作可分为数据加载 (data loading),预处理 (pre-processing),格式变化 (formatting) 和测试时数据增强 (test-time augmentation)。

下面的例子就是 PSPNet 的一个流程:

img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
crop_size = (512, 1024)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(2048, 1024), 
        #   img_scale(list或tuple):指定图像的尺度大小。
        # 可以设置为单个值
        # (如`img_scale=512`表示将图像调整为512x512的尺寸)
        # 或多个值组成的列表/元组
        # (如`img_scale=(512, 1024)`
        # 表示将图像调整为512x512和1024x1024两个尺寸)。
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,  # 指定是否进行翻转操作。默认为`True`,表示进行翻转。
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]

对于每个操作,我们列出它添加、更新、移除的相关字典域 (dict fields):(具体每个操作的含义,以及参数的作用,请结合GPT进行学习)

数据加载 Data loading

LoadImageFromFile

  • 增加: img, img_shape, ori_shape

LoadAnnotations

  • 增加: gt_semantic_seg, seg_fields

预处理 Pre-processing

Resize

  • 增加: scale, scale_idx, pad_shape, scale_factor, keep_ratio
  • 更新: img, img_shape, *seg_fields

RandomFlip

  • 增加: flip
  • 更新: img, *seg_fields

Pad

  • 增加: pad_fixed_size, pad_size_divisor
  • 更新: img, pad_shape, *seg_fields

RandomCrop

  • 更新: img, pad_shape, *seg_fields

Normalize

  • 增加: img_norm_cfg
  • 更新: img

SegRescale

  • 更新: gt_semantic_seg

PhotoMetricDistortion

  • 更新: img

格式 Formatting

ToTensor

  • 更新: 由 keys 指定

ImageToTensor

  • 更新: 由 keys 指定

Transpose

  • 更新: 由 keys 指定

ToDataContainer

  • 更新: 由 keys 指定

DefaultFormatBundle

  • 更新: img, gt_semantic_seg

Collect

  • 增加: img_meta (the keys of img_meta is specified by meta_keys)
  • 移除: all other keys except for those specified by keys

测试时数据增强 Test time augmentation

MultiScaleFlipAug

拓展和使用自定义的流程

  1. 在任何一个文件里写一个新的流程,例如 my_pipeline.py,它以一个字典作为输入并且输出一个字典

    from mmseg.datasets import PIPELINES
    
    @PIPELINES.register_module()
    class MyTransform:
    
        def __call__(self, results):
            results['dummy'] = True
            return results
    
  2. 导入一个新类

    from .my_pipeline import MyTransform
    
  3. 在配置文件里使用它

    img_norm_cfg = dict(
        mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
    crop_size = (512, 1024)
    train_pipeline = [
        dict(type='LoadImageFromFile'),
        dict(type='LoadAnnotations'),
        dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
        dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
        dict(type='RandomFlip', flip_ratio=0.5),
        dict(type='PhotoMetricDistortion'),
        dict(type='Normalize', **img_norm_cfg),
        dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
        dict(type='MyTransform'),
        dict(type='DefaultFormatBundle'),
        dict(type='Collect', keys=['img', 'gt_semantic_seg']),
    ]
    

教程 4: 自定义模型

自定义优化器 (optimizer)

假设您想增加一个新的叫 MyOptimizer 的优化器,它的参数分别为 ab, 和 c。 您首先需要在一个文件里实现这个新的优化器,例如在 mmseg/core/optimizer/my_optimizer.py 里面:

from mmcv.runner import OPTIMIZERS
from torch.optim import Optimizer


@OPTIMIZERS.register_module
class MyOptimizer(Optimizer):

    def __init__(self, a, b, c)

然后增加这个模块到 mmseg/core/optimizer/__init__.py 里面,这样注册器 (registry) 将会发现这个新的模块并添加它:

from .my_optimizer import MyOptimizer

之后您可以在配置文件的 optimizer 域里使用 MyOptimizer, 如下所示,在配置文件里,优化器被 optimizer 域所定义:

optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)

为了使用您自己的优化器,域可以被修改为:

optimizer = dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value)

我们已经支持了 PyTorch 自带的全部优化器,唯一修改的地方是在配置文件里的 optimizer 域。例如,如果您想使用 ADAM,尽管数值表现会掉点,还是可以如下修改:

optimizer = dict(type='Adam', lr=0.0003, weight_decay=0.0001)

使用者可以直接按照 PyTorch 文档教程 去设置参数。

定制优化器的构造器 (optimizer constructor)

对于优化,一些模型可能会有一些特别定义的参数,例如批归一化 (BatchNorm) 层里面的权重衰减 (weight decay)。 使用者可以通过定制优化器的构造器来微调这些细粒度的优化器参数。

from mmcv.utils import build_from_cfg

from mmcv.runner import OPTIMIZER_BUILDERS
from .cocktail_optimizer import CocktailOptimizer


@OPTIMIZER_BUILDERS.register_module
class CocktailOptimizerConstructor(object):

    def __init__(self, optimizer_cfg, paramwise_cfg=None):

    def __call__(self, model):

        return my_optimizer

开发和增加新的组件(Module)

MMSegmentation 里主要有2种组件:

  • 主干网络 (backbone): 通常是卷积网络的堆叠,来做特征提取,例如 ResNet, HRNet
  • 解码头 (decoder head): 用于语义分割图的解码的组件(得到分割结果)

添加新的主干网络

这里我们以 MobileNet 为例,展示如何增加新的主干组件:

  1. 创建一个新的文件 mmseg/models/backbones/mobilenet.py
import torch.nn as nn

from ..builder import BACKBONES


@BACKBONES.register_module
class MobileNet(nn.Module):

    def __init__(self, arg1, arg2):
        pass

    def forward(self, x):  # should return a tuple
        pass

    def init_weights(self, pretrained=None):
        pass
  1. 在 mmseg/models/backbones/__init__.py 里面导入模块
from .mobilenet import MobileNet
  1. 在您的配置文件里使用它
model = dict(
    ...
    backbone=dict(
        type='MobileNet',
        arg1=xxx,
        arg2=xxx),
    ...

增加新的解码头 (decoder head)组件

在 MMSegmentation 里面,对于所有的分割头,我们提供一个基类解码头 BaseDecodeHead 。 所有新建的解码头都应该继承它。这里我们以 PSPNet 为例, 展示如何开发和增加一个新的解码头组件:

首先,在 mmseg/models/decode_heads/psp_head.py 里添加一个新的解码头。 PSPNet 中实现了一个语义分割的解码头。为了实现一个解码头,我们只需要在新构造的解码头中实现如下的3个函数:

@HEADS.register_module()
class PSPHead(BaseDecodeHead):

    def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs):
        super(PSPHead, self).__init__(**kwargs)

    def init_weights(self):

    def forward(self, inputs):

接着,使用者需要在 mmseg/models/decode_heads/__init__.py 里面添加这个模块,这样对应的注册器 (registry) 可以查找并加载它们。

PSPNet的配置文件如下所示:

norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained='pretrain_model/resnet50_v1c_trick-2cccc1ad.pth',
    backbone=dict(
        type='ResNetV1c',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=dict(
        type='PSPHead',
        in_channels=2048,
        in_index=3,
        channels=512,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)))

增加新的损失函数

假设您想添加一个新的损失函数 MyLoss 到语义分割解码器里。 为了添加一个新的损失函数,使用者需要在 mmseg/models/losses/my_loss.py 里面去实现它。 weighted_loss 可以对计算损失时的每个样本做加权。

import torch
import torch.nn as nn

from ..builder import LOSSES
from .utils import weighted_loss

@weighted_loss
def my_loss(pred, target):
    assert pred.size() == target.size() and target.numel() > 0
    loss = torch.abs(pred - target)
    return loss

@LOSSES.register_module
class MyLoss(nn.Module):

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(MyLoss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss = self.loss_weight * my_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss

然后使用者需要在 mmseg/models/losses/__init__.py 里面添加它:

from .my_loss import MyLoss, my_loss

为了使用它,修改 loss_xxx 域。之后您需要在解码头组件里修改 loss_decode 域。 loss_weight 可以被用来对不同的损失函数做加权。

loss_decode=dict(type='MyLoss', loss_weight=1.0))

教程 5: 训练技巧

MMSegmentation 支持如下训练技巧:

主干网络和解码头组件使用不同的学习率 (Learning Rate, LR)

在语义分割里,一些方法会让解码头组件的学习率大于主干网络的学习率,这样可以获得更好的表现或更快的收敛。

在 MMSegmentation 里面,您也可以在配置文件里添加如下行来让解码头组件的学习率是主干组件的10倍。

optimizer=dict(
    paramwise_cfg = dict(
        custom_keys={
            'head': dict(lr_mult=10.)}))

通过这种修改,任何被分组到 'head' 的参数的学习率都将乘以10。您也可以参照 MMCV 文档 获取更详细的信息。
在一些常见的预训练视觉Transformer模型中,解码头组件的学习率通常设置为主干网络学习率的几倍。具体的倍数可以根据实际任务和模型的复杂程度进行调整。(一般是1-10倍),在实际应用中,可以根据实验结果进行适当的调整和优化,以获得最佳的模型性能。(来自Chat GPT3.5)

在线难样本挖掘 (Online Hard Example Mining, OHEM)(可加以应用)

对于训练时采样,我们在 这里 做了像素采样器。 如下例子是使用 PSPNet 训练并采用 OHEM 策略的配置:

_base_ = './pspnet_r50-d8_512x1024_40k_cityscapes.py'
model=dict(
    decode_head=dict(
        sampler=dict(type='OHEMPixelSampler', thresh=0.7, min_kept=100000)) )

通过这种方式,只有置信分数在0.7以下的像素值点会被拿来训练。在训练时我们至少要保留100000个像素值点。如果 thresh 并未被指定,前 min_kept 个损失的像素值点才会被选择。

类别平衡损失 (Class Balanced Loss)

对于不平衡类别分布的数据集,您也许可以改变每个类别的损失权重。这里以 cityscapes 数据集为例:

_base_ = './pspnet_r50-d8_512x1024_40k_cityscapes.py'
model=dict(
    decode_head=dict(
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0,
            # DeepLab 对 cityscapes 使用这种权重
            class_weight=[0.8373, 0.9180, 0.8660, 1.0345, 1.0166, 0.9969, 0.9754,
                        1.0489, 0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037,
                        1.0865, 1.0955, 1.0865, 1.1529, 1.0507])))

class_weight 将被作为 weight 参数,传递给 CrossEntropyLoss。详细信息请参照 PyTorch 文档 。

同时使用多种损失函数 (Multiple Losses)

对于训练时损失函数的计算,我们目前支持多个损失函数同时使用。 以 unet 使用 DRIVE 数据集训练为例, 使用 CrossEntropyLoss 和 DiceLoss 的 1:3 的加权和作为损失函数。配置文件写为:

_base_ = './fcn_unet_s5-d16_64x64_40k_drive.py'
model = dict(
    decode_head=dict(loss_decode=[dict(type='CrossEntropyLoss', loss_name='loss_ce', loss_weight=1.0),
            dict(type='DiceLoss', loss_name='loss_dice', loss_weight=3.0)]),
    auxiliary_head=dict(loss_decode=[dict(type='CrossEntropyLoss', loss_name='loss_ce',loss_weight=1.0),
            dict(type='DiceLoss', loss_name='loss_dice', loss_weight=3.0)]),
    )

通过这种方式,确定训练过程中损失函数的权重 loss_weight 和在训练日志里的名字 loss_name

注意: loss_name 的名字必须带有 loss_ 前缀,这样它才能被包括在反传的图里。

教程 6: 自定义运行设定

自定义优化设定

自定义 PyTorch 支持的优化器

我们已经支持 PyTorch 自带的所有优化器,唯一需要修改的地方是在配置文件里的 optimizer 域里面。 例如,如果您想使用 ADAM (注意如下操作可能会让模型表现下降),可以使用如下修改:

optimizer = dict(type='Adam', lr=0.0003, weight_decay=0.0001)

为了修改模型的学习率,使用者仅需要修改配置文件里 optimizer 的 lr 即可。 使用者可以参照 PyTorch 的 API 文档 直接设置参数。

自定义自己实现的优化器

1. 定义一个新的优化器

一个自定义的优化器可以按照如下去定义:

假如您想增加一个叫做 MyOptimizer 的优化器,它的参数分别有 ab, 和 c。 您需要创建一个叫 mmseg/core/optimizer 的新文件夹。 然后再在文件,即 mmseg/core/optimizer/my_optimizer.py 里面去实现这个新优化器:

from .registry import OPTIMIZERS
from torch.optim import Optimizer


@OPTIMIZERS.register_module()
class MyOptimizer(Optimizer):

    def __init__(self, a, b, c)

2. 增加优化器到注册表 (registry)

为了让上述定义的模块被框架发现,首先这个模块应该被导入到主命名空间 (main namespace) 里。 有两种方式可以实现它。

  • 修改 mmseg/core/optimizer/__init__.py 来导入它

    新的被定义的模块应该被导入到 mmseg/core/optimizer/__init__.py 这样注册表将会发现新的模块并添加它

from .my_optimizer import MyOptimizer
  • 在配置文件里使用 custom_imports 去手动导入它
custom_imports = dict(imports=['mmseg.core.optimizer.my_optimizer'], allow_failed_imports=False)

mmseg.core.optimizer.my_optimizer 模块将会在程序运行的开始被导入,并且 MyOptimizer 类将会自动注册。 需要注意只有包含 MyOptimizer 类的包 (package) 应当被导入。 而 mmseg.core.optimizer.my_optimizer.MyOptimizer 不能 被直接导入。

事实上,使用者完全可以用另一个按这样导入方法的文件夹结构,只要模块的根路径已经被添加到 PYTHONPATH 里面。

3. 在配置文件里定义优化器

之后您可以在配置文件的 optimizer 域里面使用 MyOptimizer 在配置文件里,优化器被定义在 optimizer 域里,如下所示:

optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)

为了使用您自己的优化器,这个域可以被改成:

optimizer = dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value)

自定义优化器的构造器 (constructor)

有些模型可能需要在优化器里有一些特别参数的设置,例如 批归一化层 (BatchNorm layers) 的 权重衰减 (weight decay)。 使用者可以通过自定义优化器的构造器去微调这些细粒度参数。

from mmcv.utils import build_from_cfg

from mmcv.runner.optimizer import OPTIMIZER_BUILDERS, OPTIMIZERS
from mmseg.utils import get_root_logger
from .my_optimizer import MyOptimizer


@OPTIMIZER_BUILDERS.register_module()
class MyOptimizerConstructor(object):

    def __init__(self, optimizer_cfg, paramwise_cfg=None):

    def __call__(self, model):

        return my_optimizer

默认的优化器构造器的实现可以参照 这里 ,它也可以被用作新的优化器构造器的模板。

额外的设置

优化器没有实现的一些技巧应该通过优化器构造器 (optimizer constructor) 或者钩子 (hook) 去实现,如设置基于参数的学习率 (parameter-wise learning rates)。我们列出一些常见的设置,它们可以稳定或加速模型的训练。 如果您有更多的设置,欢迎在 PR 和 issue 里面提交。

  • 使用梯度截断 (gradient clip) 去稳定训练:

    一些模型需要梯度截断去稳定训练过程,如下所示

    optimizer_config = dict(
        _delete_=True, grad_clip=dict(max_norm=35, norm_type=2))
    

    如果您的配置继承自已经设置了 optimizer_config 的基础配置 (base config),您可能需要 _delete_=True 来重写那些不需要的设置。更多细节请参照 配置文件文档 。

  • 使用动量计划表 (momentum schedule) 去加速模型收敛:

    我们支持动量计划表去让模型基于学习率修改动量,这样可能让模型收敛地更快。 动量计划表经常和学习率计划表 (LR scheduler) 一起使用,例如如下配置文件就在 3D 检测里经常使用以加速收敛。 更多细节请参考 CyclicLrUpdater 和 CyclicMomentumUpdater 的实现。

    lr_config = dict(
        policy='cyclic',
        target_ratio=(10, 1e-4),
        cyclic_times=1,
        step_ratio_up=0.4,
    )
    momentum_config = dict(
        policy='cyclic',
        target_ratio=(0.85 / 0.95, 1),
        cyclic_times=1,
        step_ratio_up=0.4,
    )
    

自定义训练计划表

我们根据默认的训练迭代步数 40k/80k 来设置学习率,这在 MMCV 里叫做 PolyLrUpdaterHook 。 我们也支持许多其他的学习率计划表:这里 ,例如 CosineAnnealing 和 Poly 计划表。下面是一些例子:

  • 步计划表 Step schedule:

    lr_config = dict(policy='step', step=[9, 10])
    
  • 余弦退火计划表 ConsineAnnealing schedule:

    lr_config = dict(
        policy='CosineAnnealing',
        warmup='linear',
        warmup_iters=1000,
        warmup_ratio=1.0 / 10,
        min_lr_ratio=1e-5)
    

自定义工作流 (workflow)

mmsegmentation/configs/base/schedules/

工作流是一个专门定义运行顺序和轮数 (running order and epochs) 的列表 (phase, epochs)。 默认情况下它设置成:

workflow = [('train', 1)]

意思是训练是跑 1 个 epoch。有时候使用者可能想检查模型在验证集上的一些指标(如 损失 loss,精确性 accuracy),我们可以这样设置工作流:

[('train', 1), ('val', 1)]

于是 1 个 epoch 训练,1 个 epoch 验证将交替运行。

注意:

  1. 模型的参数在验证的阶段不会被自动更新
  2. 配置文件里的关键词 total_epochs 仅控制训练的 epochs 数目,而不会影响验证时的工作流(0.20.2版本改为runner = dict(type='IterBasedRunner', max_iters=320000),其中max_iters为总训练的epochs数目)
  3. 工作流 [('train', 1), ('val', 1)] 和 [('train', 1)] 将不会改变 EvalHook 的行为,因为 EvalHook 被 after_train_epoch 调用而且验证的工作流仅仅影响通过调用 after_val_epoch 的钩子 (hooks)。因此, [('train', 1), ('val', 1)] 和 [('train', 1)] 的区别仅在于 runner 将在每次训练 epoch 结束后计算在验证集上的损失

自定义钩 (hooks)

mmsegmentation/configs/base/schedules、mmsegmentation/configs/base/default_runtime.py

使用 MMCV 实现的钩子 (hooks)

如果钩子已经在 MMCV 里被实现,如下所示,您可以直接修改配置文件来使用钩子:

custom_hooks = [
    dict(type='MyHook', a=a_value, b=b_value, priority='NORMAL')
]

修改默认的运行时间钩子 (runtime hooks)

以下的常用的钩子没有被 custom_hooks 注册:

  • log_config
  • checkpoint_config
  • evaluation
  • lr_config
  • optimizer_config
  • momentum_config

在这些钩子里,只有 logger hook 有 VERY_LOW 优先级,其他的优先级都是 NORMAL。 上述提及的教程已经包括了如何修改 optimizer_configmomentum_config 和 lr_config。 这里我们展示我们如何处理 log_config, checkpoint_config 和 evaluation

检查点配置文件 (Checkpoint config)

MMCV runner 将使用 checkpoint_config 去初始化 CheckpointHook.

checkpoint_config = dict(interval=1)

使用者可以设置 max_keep_ckpts 来仅保存一小部分检查点或者通过 save_optimizer 来决定是否保存优化器的状态字典 (state dict of optimizer)。 更多使用参数的细节请参考 这里 。

日志配置文件 (Log config)

log_config 包裹了许多日志钩 (logger hooks) 而且能去设置间隔 (intervals)。现在 MMCV 支持 WandbLoggerHook, MlflowLoggerHook 和 TensorboardLoggerHook。 详细的使用请参照 文档 。

log_config = dict(
    interval=50,
    hooks=[
        dict(type='TextLoggerHook'),
        dict(type='TensorboardLoggerHook')
    ])

评估配置文件 (Evaluation config)

evaluation 的配置文件将被用来初始化 EvalHook 。 除了 interval 键,其他的像 metric 这样的参数将被传递给 dataset.evaluate() 。

evaluation = dict(interval=1, metric='mIoU')

常用工具

除了训练和测试的脚本,我们在 tools/ 文件夹路径下还提供许多有用的工具。

计算参数量(params)和计算量( FLOPs) (试验性)

我们基于 flops-counter.pytorch 提供了一个用于计算给定模型参数量和计算量的脚本。

python tools/get_flops.py ${CONFIG_FILE} [--shape ${INPUT_SHAPE}]

您将得到如下的结果:

==============================
Input shape: (3, 2048, 1024)
Flops: 1429.68 GMac
Params: 48.98 M
==============================

注意: 这个工具仍然是试验性的,我们无法保证数字是正确的。您可以拿这些结果做简单的实验的对照,在写技术文档报告或者论文前您需要再次确认一下。

(1) 计算量与输入的形状有关,而参数量与输入的形状无关,默认的输入形状是 (1, 3, 1280, 800); (2) 一些运算操作,如 GN 和其他定制的运算操作没有加入到计算量的计算中。

发布模型

在您上传一个模型到云服务器之前,您需要做以下几步: (1) 将模型权重转成 CPU 张量; (2) 删除记录优化器状态 (optimizer states)的相关信息; (3) 计算检查点文件 (checkpoint file) 的哈希编码(hash id)并且将哈希编码加到文件名中。

python tools/publish_model.py ${INPUT_FILENAME} ${OUTPUT_FILENAME}

例如,

python tools/publish_model.py work_dirs/pspnet/latest.pth psp_r50_hszhao_200ep.pth

最终输出文件将是 psp_r50_512x1024_40ki_cityscapes-{hash id}.pth

导出 ONNX (试验性)

我们提供了一个脚本来导出模型到 ONNX 格式。被转换的模型可以通过工具 Netron 来可视化。除此以外,我们同样支持对 PyTorch 和 ONNX 模型的输出结果做对比。

python tools/pytorch2onnx.py \
    ${CONFIG_FILE} \
    --checkpoint ${CHECKPOINT_FILE} \
    --output-file ${ONNX_FILE} \
    --input-img ${INPUT_IMG} \
    --shape ${INPUT_SHAPE} \
    --rescale-shape ${RESCALE_SHAPE} \
    --show \
    --verify \
    --dynamic-export \
    --cfg-options \
      model.test_cfg.mode="whole"

各个参数的描述:

  • config : 模型配置文件的路径
  • --checkpoint : 模型检查点文件的路径
  • --output-file: 输出的 ONNX 模型的路径。如果没有专门指定,它默认是 tmp.onnx
  • --input-img : 用来转换和可视化的一张输入图像的路径
  • --shape: 模型的输入张量的高和宽。如果没有专门指定,它将被设置成 test_pipeline 的 img_scale
  • --rescale-shape: 改变输出的形状。设置这个值来避免 OOM,它仅在 slide 模式下可以用
  • --show: 是否打印输出模型的结构。如果没有被专门指定,它将被设置成 False
  • --verify: 是否验证一个输出模型的正确性 (correctness)。如果没有被专门指定,它将被设置成 False
  • --dynamic-export: 是否导出形状变化的输入与输出的 ONNX 模型。如果没有被专门指定,它将被设置成 False
  • --cfg-options: 更新配置选项

注意: 这个工具仍然是试验性的,目前一些自定义操作还没有被支持

评估 ONNX 模型

我们提供 tools/deploy_test.py 去评估不同后端的 ONNX 模型。

先决条件

使用方法

python tools/deploy_test.py \
    ${CONFIG_FILE} \
    ${MODEL_FILE} \
    ${BACKEND} \
    --out ${OUTPUT_FILE} \
    --eval ${EVALUATION_METRICS} \
    --show \
    --show-dir ${SHOW_DIRECTORY} \
    --cfg-options ${CFG_OPTIONS} \
    --eval-options ${EVALUATION_OPTIONS} \
    --opacity ${OPACITY} \

各个参数的描述:

  • config: 模型配置文件的路径
  • model: 被转换的模型文件的路径
  • backend: 推理的后端,可选项:onnxruntime, tensorrt
  • --out: 输出结果成 pickle 格式文件的路径
  • --format-only : 不评估直接给输出结果的格式。通常用在当您想把结果输出成一些测试服务器需要的特定格式时。如果没有被专门指定,它将被设置成 False。 注意这个参数是用 --eval 来 手动添加
  • --eval: 评估指标,取决于每个数据集的要求,例如 "mIoU" 是大多数据集的指标而 "cityscapes" 仅针对 Cityscapes 数据集。注意这个参数是用 --format-only 来 手动添加
  • --show: 是否展示结果
  • --show-dir: 涂上结果的图像被保存的文件夹的路径
  • --cfg-options: 重写配置文件里的一些设置,xxx=yyy 格式的键值对将被覆盖到配置文件里
  • --eval-options: 自定义的评估的选项, xxx=yyy 格式的键值对将成为 dataset.evaluate() 函数的参数变量
  • --opacity: 涂上结果的分割图的透明度,范围在 (0, 1] 之间

结果和模型

模型配置文件数据集评价指标PyTorchONNXRuntimeTensorRT-fp32TensorRT-fp16
FCNfcn_r50-d8_512x1024_40k_cityscapes.pycityscapesmIoU72.272.272.272.2
PSPNetpspnet_r50-d8_512x1024_40k_cityscapes.pycityscapesmIoU77.877.877.877.8
deeplabv3deeplabv3_r50-d8_512x1024_40k_cityscapes.pycityscapesmIoU79.079.079.079.0
deeplabv3+deeplabv3plus_r50-d8_512x1024_40k_cityscapes.pycityscapesmIoU79.679.579.579.5
PSPNetpspnet_r50-d8_769x769_40k_cityscapes.pycityscapesmIoU78.278.1
deeplabv3deeplabv3_r50-d8_769x769_40k_cityscapes.pycityscapesmIoU78.578.3
deeplabv3+deeplabv3plus_r50-d8_769x769_40k_cityscapes.pycityscapesmIoU78.9

注意: TensorRT 仅在使用 whole mode 测试模式时的配置文件里可用。

导出 TorchScript (试验性)

我们同样提供一个脚本去把模型导出成 TorchScript 格式。您可以使用 pytorch C++ API LibTorch 去推理训练好的模型。 被转换的模型能被像 Netron 的工具来可视化。此外,我们还支持 PyTorch 和 TorchScript 模型的输出结果的比较。

python tools/pytorch2torchscript.py \
    ${CONFIG_FILE} \
    --checkpoint ${CHECKPOINT_FILE} \
    --output-file ${ONNX_FILE}
    --shape ${INPUT_SHAPE}
    --verify \
    --show

各个参数的描述:

  • config : pytorch 模型的配置文件的路径
  • --checkpoint : pytorch 模型的检查点文件的路径
  • --output-file: TorchScript 模型输出的路径,如果没有被专门指定,它将被设置成 tmp.pt
  • --input-img : 用来转换和可视化的输入图像的路径
  • --shape: 模型的输入张量的宽和高。如果没有被专门指定,它将被设置成 512 512
  • --show: 是否打印输出模型的追踪图 (traced graph),如果没有被专门指定,它将被设置成 False
  • --verify: 是否验证一个输出模型的正确性 (correctness),如果没有被专门指定,它将被设置成 False

注意: 目前仅支持 PyTorch>=1.8.0 版本

注意: 这个工具仍然是试验性的,一些自定义操作符目前还不被支持

例子:

  • 导出 PSPNet 在 cityscapes 数据集上的 pytorch 模型

    python tools/pytorch2torchscript.py configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py \
    --checkpoint checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
    --output-file checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pt \
    --shape 512 1024
    

导出 TensorRT (试验性)

一个导出 ONNX 模型成 TensorRT 格式的脚本

先决条件

使用方法

python ${MMSEG_PATH}/tools/onnx2tensorrt.py \
    ${CFG_PATH} \
    ${ONNX_PATH} \
    --trt-file ${OUTPUT_TRT_PATH} \
    --min-shape ${MIN_SHAPE} \
    --max-shape ${MAX_SHAPE} \
    --input-img ${INPUT_IMG} \
    --show \
    --verify

各个参数的描述:

  • config : 模型的配置文件
  • model : 输入的 ONNX 模型的路径
  • --trt-file : 输出的 TensorRT 引擎的路径
  • --max-shape : 模型的输入的最大形状
  • --min-shape : 模型的输入的最小形状
  • --fp16 : 做 fp16 模型转换
  • --workspace-size : 在 GiB 里的最大工作空间大小 (Max workspace size)
  • --input-img : 用来可视化的图像
  • --show : 做结果的可视化
  • --dataset : Palette provider, 默认为 CityscapesDataset
  • --verify : 验证 ONNXRuntime 和 TensorRT 的输出
  • --verbose : 当创建 TensorRT 引擎时,是否详细做信息日志。默认为 False

注意: 仅在全图测试模式 (whole mode) 下测试过

其他内容

打印完整的配置文件

tools/print_config.py 会逐字逐句的打印整个配置文件,展开所有的导入。

python tools/print_config.py \
  ${CONFIG} \
  --graph \
  --cfg-options ${OPTIONS [OPTIONS...]} \

各个参数的描述:

  • config : pytorch 模型的配置文件的路径
  • --graph : 是否打印模型的图 (models graph)
  • --cfg-options: 自定义替换配置文件的选项

对训练日志 (training logs) 画图

tools/analyze_logs.py 会画出给定的训练日志文件的 loss/mIoU 曲线,首先需要 pip install seaborn 安装依赖包。

python tools/analyze_logs.py xxx.log.json [--keys ${KEYS}] [--legend ${LEGEND}] [--backend ${BACKEND}] [--style ${STYLE}] [--out ${OUT_FILE}]

示例:

  • 对 mIoU, mAcc, aAcc 指标画图

    python tools/analyze_logs.py log.json --keys mIoU mAcc aAcc --legend mIoU mAcc aAcc
    
  • 对 loss 指标画图

    python tools/analyze_logs.py log.json --keys loss --legend loss
    

转换其他仓库的权重

tools/model_converters/ 提供了若干个预训练权重转换脚本,支持将其他仓库的预训练权重的 key 转换为与 MMSegmentation 相匹配的 key。

ViT Swin MiT Transformer 模型

  • ViT

tools/model_converters/vit2mmseg.py 将 timm 预训练模型转换到 MMSegmentation。

python tools/model_converters/vit2mmseg.py ${SRC} ${DST}
  • Swin

    tools/model_converters/swin2mmseg.py 将官方预训练模型转换到 MMSegmentation。

    python tools/model_converters/swin2mmseg.py ${SRC} ${DST}
    
  • SegFormer

    tools/model_converters/mit2mmseg.py 将官方预训练模型转换到 MMSegmentation。

    python tools/model_converters/mit2mmseg.py ${SRC} ${DST}
    

模型服务

为了用 TorchServe 服务 MMSegmentation 的模型 , 您可以遵循如下流程:

1. 将 model 从 MMSegmentation 转换到 TorchServe

python tools/mmseg2torchserve.py ${CONFIG_FILE} ${CHECKPOINT_FILE} \
--output-folder ${MODEL_STORE} \
--model-name ${MODEL_NAME}

注意: ${MODEL_STORE} 需要设置为某个文件夹的绝对路径

2. 构建 mmseg-serve 容器镜像 (docker image)

docker build -t mmseg-serve:latest docker/serve/

3. 运行 mmseg-serve

请查阅官方文档: 使用容器运行 TorchServe

为了在 GPU 环境下使用, 您需要安装 nvidia-docker. 若在 CPU 环境下使用,您可以忽略添加 --gpus 参数。

示例:

docker run --rm \
--cpus 8 \
--gpus device=0 \
-p8080:8080 -p8081:8081 -p8082:8082 \
--mount type=bind,source=$MODEL_STORE,target=/home/model-server/model-store \
mmseg-serve:latest

阅读关于推理 (8080), 管理 (8081) 和指标 (8082) APIs 的 文档 。

4. 测试部署

curl -O https://raw.githubusercontent.com/open-mmlab/mmsegmentation/master/resources/3dogs.jpg
curl http://127.0.0.1:8080/predictions/${MODEL_NAME} -T 3dogs.jpg -o 3dogs_mask.png

得到的响应将是一个 ".png" 的分割掩码.

您可以按照如下方法可视化输出:

import matplotlib.pyplot as plt
import mmcv
plt.imshow(mmcv.imread("3dogs_mask.png", "grayscale"))
plt.show()

看到的东西将会和下图类似: image.png