MMSegmentation 是 OpenMMLab 开源的基于 PyTorch 实现的功能强大的语义分割工具箱 \
MMSegmentation 的主要特性如下:
- 丰富的语义分割模型: 已支持 11 种主干网络和 34 种算法,例如常用模型 FCN, PSPNet 和 DeepLabV3;Transformer 模型,Swin Transformer、Segmenter 和 SegFormer; Real-Time 实时分割模型, ICNet、BiSeNet 和 STDC 等;以及最近流行的网络 ConvNeXt 和 MAE。(具体支持哪些网络,可查看官方文档)
- 大量开箱即用的模型权重:在 16 个常用的语义分割数据集上提供了 590 个训练好的模型。
- 统一的性能评估框架:优化和统一了训练和测试的流程,方便公平比较各个模型在特定任务上的表现。
目前,MMSegmentation 支持的分割任务为语义分割 ,MMDetection 中支持了实例分割和全景分割。
MMSegmentation 算法库框架介绍
本文基于MMSegmentation v0.23.0 的整体架构(2022.05.26)
具体当前已实现的框架得结合当前最新代码及文档!
MMSegmentation 目录结构
按照代码目录下的文件夹,MMSegmentation 代码库主要可以包含四个部分:
1)./tools 包括了调用 MMSegmentation 作为训练和测试入口的./tools/train.py 和 ./tools/test.py,预训练模型和数据集准备的转换脚本,以及部署和可视化相关的脚本。
详细介绍可见 Github 里的文档
2)./configs 包括了各个算法的配置文件、存放常用的数据集配置、基础模型以及训练策略的基配置文件 ./configs/_base_。
3)./mmseg 里面是 MMSegmentation 的算法库,包括核心组件、数据集处理、分割模型代码和面向用户的 API 接口。
4)./data 指的是存放数据集的路径,在原本的代码库中没有这个文件夹。用户只需指定正确的文件夹路径即可使用数据。
下面是详细的 MMSegmentation 的算法库目录结构:(具体可直接看代码)
# MMSegmentation 算法库目录结构的主要部分
mmsegmentation
|
|- configs #配置文件
|- _base_ ## 基配置文件
| | |- datasets ### 数据集相关配置文件
| | |- models ### 模型相关配置文件
| | |- schedules ### 训练日程如优化器,学习率等相关配置文件
| | |- default_runtime.py ### 运行相关的默认的设置
| |- swin ## 各个分割模型的配置文件,会引用 _base_ 的配置并做修改
| |- ...
|- data # 原始及转换后的数据集文件
|- mmseg
| |- core ## 核心组件
| | |- evaluation ### 评估模型性能代码
| |- datasets ## 数据集相关代码
| | |- pipelines ### 数据预处理
| | |- samplers ### 数据集采样代码
| | |- ade.py ### 各个数据集准备需要的代码
| | |- ...
| |- models ## 分割模型具体实现代码
| | |- backbones ### 主干网络
| | |- decode_heads ### 解码头
| | |- losses ### 损失函数
| | |- necks ### 颈
| | |- segmentors ### 构建完整分割网络的代码
| | |- utils ### 构建模型时的辅助工具
| |- apis ## high level 用户接口,在这里调用 ./mmseg/ 内各个组件
| | |- train.py ### 训练接口
| | |- test.py ### 测试接口
| | |- ...
| |- ops ## cuda 算子(即将迁移到 mmcv 中)
| |- utils ## 辅助工具
|- tools
| |- model_converters ## 各个主干网络预训练模型转 key 脚本
| |- convert_datasets ## 各个数据集准备转换脚本
| |- train.py ## 训练脚本
| |- test.py ## 测试脚本
| |- ...
|- ...
MMSegmentation 的算法库有 3 个关键组件:
1)./mmseg/apis/,用于训练和测试的接口。
2)./mmseg/models/,用于分割网络模型的具体实现。
3)./mmseg/datasets/,用于数据集处理。
本文我们主要介绍网络模型结构,因此涉及内容主要在 ./mmseg/models 里面。
MMSegmentation 模型实现
Segmentor
MMSegmentation 中将语义分割模型定义为 segmentor, 一般包括 backbone、neck、head、loss 4 个核心组件,每个模块的功能如下:
- 预处理后的数据输入到 backbone(如 ResNet 和 Swin Transformer)中进行编码并提取特征。
- 输出的单尺度或者多尺度特征图输入到 neck 模块中进行特征融合或者增强,典型的 neck 是 特征金字塔(Feature Pyramid Networks, FPN)。
- 上述多尺度特征最终输入到 head 部分,一般包括 decoder head, auxiliary head 以及 cascade decoder head,用以预测分割结果(它们的区别我们会在下文具体介绍)。
- 最后一步是计算 pixel 分类的 loss,进行训练。
需要说明的是,上述 4 个组件不是每个算法都需要的,比如很多模型里没有 neck 和 auxiliary head 组件。分割器 segmentor 的具体代码见文件 ./mmseg/models/segmentors/。
MMSegmentation 里面的分割器框架可以分为 “Encoder Decoder” 结构和 “Cascade Encoder Decoder” 结构。现有的大多数模型为 “Encoder Decoder” 结构,即利用 encoder 提取图像特征,再用 decoder 去解码上述特征。 “Cascade Encoder Decoder” 的解码部分不是单独的解码头,而是级联式的 2 个或多个解码头,前一个解码头的输出作为后一个解码头的输入。
关于分割器 segmentor 的训练和测试的基本逻辑,以语义分割经典的 “Encoder Decoder” 结构为例:
class EncoderDecoder(BaseSegmentor):
def __init__(...):
# 构建 backbone、neck 和 head
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
self._init_decode_head(decode_head)
self._init_auxiliary_head(auxiliary_head)
def forward_train(...):
# 利用 backbone+neck 进行特征提取
x = self.extract_feat(img)
losses = dict()
# decode head 输出预测特征图并计算出 loss
loss_decode = self._decode_head_forward_train(x, img_metas, gt_semantic_seg)
losses.update(loss_decode)
# auxiliary heads 输出预测特征图并计算出 loss
if self.with_auxiliary_head:
loss_aux = self._auxiliary_head_forward_train(x, img_metas, gt_semantic_seg)
losses.update(loss_aux)
return losses
def simple_test(...):
# 调用 inference 函数,对输入图片做全图或者滑动窗口的推理,得到 logits 值
seg_logit = self.inference(img, img_meta, rescale)
# 做 argmax 得到预测的 prediction mask
seg_pred = seg_logit.argmax(dim=1)
def aug_test(...):
...
接下来,我们详细介绍分割器 segmentor 里 4 个核心组件:backbone, neck,head,和 loss。
Backbone
目前 MMSegmengtation 中已经集成了大部分主干网络,具体见文件 ./mmseg/models/backbones/,v0.24.1 已经实现的骨架如下:
通常定义的“主干网络” 是指从上游任务(如 ImageNet )预训练,然后用于多个下游任务(如目标检测、实例分割、语义分割、姿态估计)中的网络,而在
./mmseg/models/backbones 里主干网络的定义有所不同,会把一些分割算法的网络结构也作为“主干网络”,如 UNet、 FastSCNN、CGNet、ICNet、BiSeNetV1/V2、ERFNet、STDC。
其中最常用的是 ResNet v1c 系列和 Vision Transformer 系列。如果你需要对骨架进行扩展,可以继承上述网络,然后通过注册器机制注册使用。一个典型用法为 ./configs/_base_/models/segmenter_vit-b16_mask.py 里面的:
checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_base_p16_384_20220308-96dfe169.pth' # noqa#
model settings
model = dict(
type='EncoderDecoder',
pretrained=checkpoint, # 加载的预训练模型,这里为 Google Research提供的由 JAX 训练框架得到的 Vision Transformer
backbone=dict(
type='VisionTransformer', # 骨架类名,后面的参数都是该类的初始化参数
img_size=(512, 512),
patch_size=16,
in_channels=3,
embed_dims=768,
num_layers=12,
num_heads=12,
drop_path_rate=0.1,
attn_drop_rate=0.0,
drop_rate=0.0,
final_norm=True,
norm_cfg= dict(type='LN', eps=1e-6, requires_grad=True),
with_cls_token=True,
interpolate_mode='bicubic',
),
同 OpenMMLab 其他算法库一样,我们使用了 MMCV 中的模块注册机制,通过修改配置文件的 type ,可以使用在 MMSegmentation 已经实现的 backbone 模型。此外,还可以使用 MMClassification 里面的更多主干网络,如 ShuffleNet、EfficientNet 等,可根据 ./configs/convnext 里面 ConvNeXt 的实现方式,详细的方式可以参考: OpenMMLab 模型大联动,MMDet 也能用 MMCls 的网络。
MMCV 模块注册机制网址
Neck
Neck 可以认为是 backbone 和 head 的连接层,主要负责对 backbone 的特征进行高效融合和增强,能够对输入的单尺度或者多尺度特征进行融合、增强输出等。具体见文件 ./mmseg/models/necks/,v0.24.1 已经实现的 neck 如下:
最常用的应该是 FPN,一个典型用法是
./configs/_base_/models/pointrend_r50.py 里面:
neck=dict(
type='FPN',
in_channels=[256, 512, 1024, 2048], # 骨架多尺度特征图输出通道
out_channels=256, # 增强后通道输出
num_outs=4), # 输出num_outs个多尺度特征图
Head
MMSegmentation 的 head 是用来处理 backbone 或 neck 的特征图,对图像里的每个像素 pixel 做分类然后得到分类的结果。具体见文件 ./mmseg/models/decode_heads/,v0.24.1 已经实现的 head 如下:
虽然它们都是用来解码特征图中的信息。但在使用上,可以将它们分为 decoder head,auxiliary head 以及 cascade decoder head。
- decoder head 是直接在训练和推理中作为图像预测输出的 head。
- auxiliary head 是只在训练过程中输出图像预测用来辅助损失函数计算的 head。
- cascade decoder head 是指级联式的 2 个或多个解码头,前一个解码头的输出作为后一个解码头的输入, OCRNet 和 PointRend 两种算法就使用了 cascade decoder head。
在 MMSegmentation 里每个 head 自己单独计算损失,所以把这个公共的行为抽象成了一个基类: BaseDecodeHead,每个算法的 head 都继承自这个基类,类里面包括了计算 loss 的函数,用于计算 head 输出的 logits 值和 label 的损失。
Loss
MMSegmentation 里的 loss 计算的是每个像素上的 logits 和分割标签之间的差别,使用最多的是 cross entropy loss 和 dice loss,v0.24.1 已经实现的 loss 如下:
除了
./mmseg/models/losses/ 里的这些 loss 外,计算 loss 时还可以用到一些策略和方法,比如:在线难样本挖掘策略 (OHEM, Online Hard Example Mining) 。
MMSegmentation 跑语义分割数据集
2022.06.07
数据集
无论是MMSegmentation已经支持的数据集,还是自己的数据集,都需要在配置文件里配置数据相关的信息,如数据集本地存储路径,数据预处理流程 Pipeline;继承数据集基类 CustomDataset 以方便调用在某个数据集上加载图像和标注,解析加载数据,评估模型表现等各种功能。
数据配置文件
MMSegmentation 的数据集配置基文件在 ./configs/base/datasets 里面,每个数据集配置文件主要包括:(1)data 字段,主要包括dataloader 的配置,例如模型训练时每个 GPU 上面的样本数目和进程数;(2)数据集和数据预处理配置,例如数据集路径和数据预处理 Pipeline。
数据配置文件的 data 字段
这是数据配置文件的一个样例:
data = dict(
samples_per_gpu=4,
workers_per_gpu=4,
train=dict(
type='ADE20KDataset',
data_root='data/ade/ADEChallengeData2016',
img_dir='images/training',
ann_dir='annotations/training',
pipeline=train_pipeline),
val=dict(
type='ADE20KDataset',
data_root='data/ade/ADEChallengeData2016',
img_dir='images/validation',
ann_dir='annotations/validation',
pipeline=test_pipeline),
test=dict(
type='ADE20KDataset',
data_root='data/ade/ADEChallengeData2016',
img_dir='images/validation',
ann_dir='annotations/validation',
pipeline=test_pipeline))
data 中重要的是如下几个字段:
train,valandtest: 构建数据集实例的配置,可以通过registry&build机制来构建,分别用于模型的训练、验证和测试。samples_per_gpu: 在模型训练时每个 GPU 加载的样本数,它乘以模型训练时的 GPU 数目就是模型训练时的batch_size。例如,当使用 8 块 GPU 做分布式训练并且samples_per_gpu=4,那么batch_sizeis8*4=32。如果想定义不同batch_size用于验证和测试,需要在版本 >=0.24.1 的 MMSegmentation 中使用val_dataloader和test_dataloaser。workers_per_gpu: 数据加载时每个 GPU 使用的子进程(subprocess)数目。0则意味着主进程加载数据。
需要说明的是,samples_per_gpu 仅用于模型训练,因为目前 MMSegmentation 并不支持 batch 方式的推理,所以验证和测试时 samples_per_gpu=1,即每张 GPU 的样本数都是 1。
MMSegmentation 在 v0.24.1 之前,除了 train、 val、test、samples_per_gpu 和 workers_per_gpu,data 中的其他字段必须是 PyTorch 中 dataloader 的输入参数,并且模型训练、验证和测试的 dataloaders 都有着同样的输入参数。在 v0.24.1 之后,尽管上述的参数定义仍然可用,但是会将优先支持使用 train_dataloader、val_dataloader和 test_dataloaser 去分别指定模型训练、验证和测试时 dataloader 所需要的参数。
以下就是一个train_dataloader、val_dataloader和 test_dataloaser 使用不同参数的样例:
data = dict(
samples_per_gpu=4,
workers_per_gpu=4,
shuffle=True,
train=dict(type='xxx', ...),
val=dict(type='xxx', ...),
test=dict(type='xxx', ...),
# 在验证和测试时使用不同的 batch size
val_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False),
test_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False))
假如只有一张 GPU 用于模型的训练和测试,因为整体 dataloader 参数定义的优先级比较低,所以训练的 batch size 是 4 并且数据集将会被 shuffle,验证和测试的 batch size 是 1 并且数据集不会被 shuffle。
在 MMSegmentation v0.24.1 之后,我们更推荐使用专门的 dataloader 设置去替代整体 dataloader 的定义,这样可以让数据配置更加清晰易懂。可以修改为:
data = dict(
train=dict(type='xxx', ...),
val=dict(type='xxx', ...),
test=dict(type='xxx', ...),
# 使用特定的 dataloader 设置
train_dataloader=dict(samples_per_gpu=4, workers_per_gpu=4, shuffle=True),
val_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False),
test_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False))
数据集预处理 Pipeline
Pipeline 由一系列数据预处理模块组成,得益于 MM 系列模块化的特性,每个模块也都可以单独配置并通过 registry&build 机制来构建。
上图是 MMSegmentation 典型的训练流程 Pipeline,每个模块都接收字典输入,输出也是字典。按照从左到右的顺序执行,绿色表示该模块运行后的新增字段,橙色表示该模块运行后被修改的字段。Pipeline 中的数据变换可以被划分如下:
- 图片和标签加载,例如
LoadImageFromFile和LoadAnnotations。 - 数据处理,例如
RandomFlip,PhotoMetricDistortion和Resize等,这部分是通常在训练流程中使用。 - 数据收集,例如
Collect,它会重新收集数据处理的字典,用来准备输入到模型里的数据
在 MMSegmentation 框架中,图片和标签加载和数据处理流程一般是固定的,用户在定制自己数据集的时候,也需要结合具体情况构建合适的 Pipeline。以 ADE20K 数据集为例,在配置文件中,训练时的 Pipeline 如下所示:
# dataset settings
dataset_type = 'ADE20KDataset' # 数据集类型,这将被用来定义数据集。
data_root = 'data/ade/ADEChallengeData2016' # 数据的根路径。
img_norm_cfg = dict( # 图像归一化配置,用来归一化输入的图像。
mean=[123.675, 116.28, 103.53], # 预训练里用于预训练主干网络模型的平均值。
std=[58.395, 57.12, 57.375], # 预训练里用于预训练主干网络模型的标准差。
to_rgb=True) # 预训练里用于预训练主干网络的图像的通道顺序。
crop_size = (512, 512) # 训练时的裁剪大小
train_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='LoadAnnotations', reduce_zero_label=True),
dict(type='Resize', # 变化图像和其注释大小的数据增广。
img_scale=(2048, 512), # 图像和标注的 resize 尺度
ratio_range=(0.5, 2.0)), # 随机 resize 的比例范围。
dict(type='RandomCrop', # 随机裁剪当前图像和其注释。
crop_size=crop_size, # 随机裁剪图像生成 patch 的大小。
cat_max_ratio=0.75), # 单个类别可以填充的最大区域的比例。
dict(type='RandomFlip', # 翻转图像和其注释。
prob=0.5), # 翻转图像的概率
dict(type='PhotoMetricDistortion'), # 光学上使用一些方法扭曲当前图像。
dict(type='Normalize', # 归一化当前图像数据。
**img_norm_cfg),
dict(type='Pad', # 填充当前图像到指定大小。
size=crop_size, # 填充的图像大小。
pad_val=0, # 图像的填充值。
seg_pad_val=255), # 'gt_semantic_seg'的填充值。
dict(type='DefaultFormatBundle'), # 默认格式转换的组合操作。
dict(type='Collect', keys=['img', 'gt_semantic_seg']), # 决定数据里哪些键被传递到分割器里的流程。
]
在构建自己的 Pipeline 时,一定要仔细检查是否真正调用了修改的配置文件,因为新增和修改的字典一旦被错误地覆盖或者遗漏,在运行程序时也可能不会报错,使得排查错误变得困难。
CustomDataset介绍
在训练或验证时经常需要获取该数据集的相关信息,比如获取数据集注释的相关信息,评估数据集预测结果的某些评价指标等等。因此,在 ./mmseg/datasets/custom.py 里面把数据集抽象成一个基类 CustomDataset,在基类里面定义了这些基本的函数,以方便被调用。
CustomDataset 里面主要有以下几个函数:
-
load_annotations():加载全部标注文件,返回一个 List, 其中每个元素是一个字典,键分别是filename和ann,值里的信息是图片和对应的标注的文件名。 -
get_ann_info(idx):根据输入inx获取对应语义分割标注的文件名。 -
prepare_train_img(idx):获得经过训练数据处理 Pipeline 之后的训练集的图像数组img,和对应的元信息img_metas。img_metas里的内容可见上一章介绍 Pipeline 的示意图。 -
prepare_test_img(idx):获得经过测试数据流 Pipeline 之后的验证集或测试集的图像数组,和其对应的元信息。字典里内容和prepare_train_img()的一样。 -
__getitem__(idx):通过判断当前是否为训练模式来调用prepare_train_img(idx)或prepare_test_img(idx)。 -
evaluate(): 评估数据集,输入预测的结果,返回这个数据集所需要的一些评价指标。
因为MMSegmentation 数据集都继承自 CustomDataset,所以熟悉它便熟悉了MMSegmentation 其他数据集的加载、解析和评估的流程。
在介绍完数据集配置文件中需要加入的预处理 Pipeline 和数据集需要继承的 CustomDataset 类之后,下面介绍如何处理自己的数据集,以便训练或验证。
数据集准备和定制化
数据集的准备和定制化,具体可以分成以下几步:
-
数据集准备,推荐在 mmsegmetaion 目录新建路径 data,然后将数据集转换成 MMSegmentation 可用的格式:分别定义好数据集图像和标注的正确文件路径,其中的标注格式为仅包含每个像素对应标签 id 的单通道标注文件,而不是三通道的 RGB 格式。
-
在
./mmseg/datasets里定义该数据集以注册到DATASETS里。 -
在
./configs/_base_/datasets里面设置训练与验证时数据集配置的参数,如数据集路径,数据增强策略等。
上述步骤主要改动的文件位置为:
mmsegmentation
|
|- data
| |- my_dataset # 转换后的自己的数据集文件
|- mmseg
| |- datasets
| | |- __init__.py # 在这里加入自己的数据集的类
| | |- my_dataset.py ## 定义自己的数据集的类
| | |- ...
|- configs
| |- _base_
| | |- datasets
| | | |- my_dataset_config.py # 自己的数据集的配置文件
| | |- ...
| |- ...
接下来我们详细介绍这三步。
数据集准备
在使用模型做训练、验证和推理前,需要将数据集处理成 MMSegmentation 定制化的格式。对于 MMSegmentation 已经支持的数据集,我们在 ./tools/convert_datasets 中提供了数据集的转换脚本,它们会转换这些原始数据集的大小(例如将较大的遥感数据集裁剪成较小的)和内容(例如将 RBG 格式的标注转换成仅包含每个像素对应标签 id 的单通道标注),同时改变图像和标注的格式与文件夹结构。相关命令可参考数据集准备文档,转换后的数据集可以被 MMSegmentation 一键运行。
数据集最终的目录组织如下,需要将图片放到 img_dir 下,对应的分割标注放到 ann_dir 下:
|- 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_suffix} 和 {seg_map_suffix} 是图像和标注的后缀,常用的是 .png 和 .jpg。
实现自己的数据集
生成好上述数据格式后,在 ./mmseg/dataset 里实现数据集,使它可以被注册到 MMCV 的 DATASETS 里面然后被模型调用。实现自己的数据集,只需要继承 CustomDataset 这个类,再定义数据集标注的名称、可视化调色盘以及文件夹后缀格式,如下所示:
from .builder import DATASETS
from .custom import CustomDataset
#将 MyDataset 类注册到 DATASETS 里
@DATASETS.register_module()
class MyDataset(CustomDataset):
# 数据集标注的各类名称,即 0, 1, 2, 3... 各个类别的对应名称
CLASSES = ('label_a', 'label_b', 'label_c', 'label_d',
'label_e', ...)
# 各类类别的 BGR 三通道值,用于可视化预测结果
PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0],
[255, 255, 0], ...]
# 图片和对应的标注,这里对应的文件夹下均为 .png 后缀
def __init__(self, **kwargs):
super(MyDataset, self).__init__(
img_suffix='.png',
seg_map_suffix='.png',
reduce_zero_label=False, # 此时 label 里的 0(上面 CLASSES 里第一个 “label_a”)在计算损失函数和指标时不会被忽略。
**kwargs)
在 ./mmseg/dataset/my_dataset.py 里面定义了数据集的分割类别 CLASSES 和对应的 BGR 通道的调色板 PALETTE,PALETTE 只在预测结果可视化的时候会用到,并不会影响训练和验证。需要强调的是,如果 label 中的 0 是背景并且想在计算评价指标的时候忽略掉它,需要设置 reduce_zero_label=True。
其在 ./mmseg/core/evaluation/metrics.py 中的原理是:当设置 reduce_zero_label=True 时,会修改分割的标签类别,将 index 为 0 的类别安排到 255,所以在训练和和测试加载分割标注时,都会做如下操作:
if reduce_zero_label:
label[label == 0] = 255
label = label - 1
label[label == 254] = 255
255 是标签里被忽略的 index。创建好 ./mmseg/dataset/my_dataset.py 后,需要在 ./mmseg/dataset/__init__.py 里也加入它:
# Copyright (c) OpenMMLab. All rights reserved.
from .my_dataset import MyDataset
__all__ = [
...,
'MyDataset'
]
设置数据集配置文件
数据集定义好后,还需要在 ./configs/_base_/datasets 里面定义该数据集有关的配置项 my_dataset_config.py,使之与其他的配置参数一起在训练和测试时调用。
首先简单介绍下为何多了一个 _base_ 基配置文件夹:自从 2020 年 6 月 MMDetection 发布 V2.0 版本以来,OpenMMLab 代码库设计了新的 config 系统,支持了多重继承机制。将常用的数据集配置、基础模型以及训练策略放到了 ./configs/_base_/ 文件夹中。每个新的 config 只需要继承一个或者多个已有的 config,然后对其中需要修改的字段进行重载。通过将 config 继承的层级控制在可接受范围内,提升了配置文件的可维护性。
以下为数据集配置文件的一个示例:
# 在./mmseg/datasets/__init__.py 中定义的数据集类型
dataset_type = 'MyDataset'
# 数据集准备生成的文件夹路径
data_root = 'data/my_dataset'
img_norm_cfg = dict( # 常用这组参数归一化是因为它是 ImageNet 1K 预训练使用的图像均值与方差
mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375],
to_rgb=True
)
crop_size = (512, 512) # 训练时图像裁剪的大小
train_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='LoadAnnotations', reduce_zero_label=True),
dict(type='Resize', img_scale=(512, 512), ratio_range=(0.5, 2.0)),
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
dict(type='RandomFlip', prob=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=(512, 512),
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
flip=False,
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']),
])
]
data = dict(
samples_per_gpu=4, # 单个 GPU 的 Batch size
workers_per_gpu=4, # 单个 GPU 分配的数据加载线程数
train=dict( # 训练数据集配置
type=dataset_type, # 数据集的类别, 细节参考自 mmseg/datasets/
data_root=data_root, # 数据集的根目录。
img_dir='img_dir/train', # 数据集图像的文件夹
ann_dir='ann_dir/train', # 数据集注释的文件夹
pipeline=train_pipeline), # 流程, 由之前创建的 train_pipeline 传递进来
val=dict( # 验证数据集的配置
type=dataset_type,
data_root=data_root,
img_dir='img_dir/val',
ann_dir='ann_dir/val',
pipeline=test_pipeline), # 由之前创建的 test_pipeline 传递的流程
test=dict(
type=dataset_type,
data_root=data_root,
img_dir='img_dir/val',
ann_dir='ann_dir/val',
pipeline=test_pipeline)
)
各个配置项的具体作用可以参考配置文件教程。至此,定义的数据集就完成了数据集的准备和定制化,只需要在 ./configs/ 里创建的配置文件里调用该数据集即可。例如:
_base_ = [
'../_base_/models/pspnet_r50-d8.py', '../_base_/datasets/my_dataset_config.py',
'../_base_/default_runtime.py', '../_base_/schedules/schedule_80k.py'
]
model = dict(
decode_head=dict(num_classes=YOUR_DATASET_CLASSES), auxiliary_head=dict(num_classes=YOUR_DATASET_CLASSES))
本文主要参考公众号OpenMMLab的文章《带你轻松掌握 MMSegmentation 整体构建流程》、《# 带你轻松用 MMSegmentation 跑语义分割数据集》