mmsegmentation二分类配置

Author Avatar
patrickcty 5月 21, 2021

教训

configs 目录下有很多已有的配置,在自己乱试之前先看看!比如医学图像分割也是二分类任务,并且其最有名的模型是 U-Net,因此今天的配置都可以直接参考相关的配置。

配置环境

由于 mmsegmentation 的版本更新得比较快,因此基于此开发的项目环境一般不太兼容,最好每个都重新搞一个虚拟环境,命令如下:

conda create -n open-mmlab python=3.7 -y
conda activate open-mmlab

conda install pytorch=1.6.0 torchvision cudatoolkit=10.1 -c pytorch -y
pip install mmcv-full==1.2.2 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.6.0/index.html
git clone https://github.com/fudan-zvg/SETR.git  # 某个基于 mmsegmentation 开发的项目
cd SETR
pip install -e .  # or "python setup.py develop"
pip install -r requirements/optional.txt

# 可选,这里是为了切换多个 CUDA 版本
conda env config vars set PATH=/home/sse/anaconda3/envs/open-mmlab/bin:/usr/local/cuda-10.1/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-10.1/lib64:/usr/local/cuda-10.1/extras/CUPTI/lib64:$LD_LIBRARY_PATH -n open-mmlab

纲要

  • 自定义数据集
    • configs/_base_/datasets/<your_dataset>.py
    • mmseg/datasets/<your_dataset>.py
    • mmseg/datasets/__init__.py
  • 自定义训练

自定义数据集

MM 系列框架都经过了非常高的封装,我们只用修改配置文件就可以非常容易地训练不同的模型。其本质上就是通过各个模块的注册来实现的。比如说在 mmseg/datasets 目录下有一个 builder.py 文件,这个文件中有 DATASETSPIPELINE 两个全局变量,前者用来注册数据集,后者用来注册数据加载处理的一些操作。比如打开 ade.py 这个数据集定义文件,我们可以看到在定义类的时候有一个装饰器 `@DATASETS.register_module()自动注册了这个数据集。因此如果要自定义数据集甚至是网络的模块,只用在相关位置定义然后注册即可。注册的话还有另一个步骤,那就是将其添加到mmseg/datasets/init.py` 之中。

创建数据集类

先别慌着写,先把已有的数据集类都扫一遍。扫一遍之后我们可以发现要自定义一个数据集类,首先需要继承 CustomDataset 类,然后再在其中自定义 CLASSES 类变量与 PALETTE 类变量,其中前者是为了将类别标号与自然语言对应起来,后者是可视化的时候对应不同的颜色。再就只用根据数据集的特点来自定义构造方法了。我们先看看 CustomDataset 构造方法的参数,在 mmseg/datasets/custom.py 的 47 行,由于其中大部分都通过配置文件来指定,因此我们只集中于几个重要的参数:

Args:
    img_suffix (str): 图像的后缀,只读取特定后缀的图像,对于混合搭建的数据集非常友好,默认值:'.jpg'
    seg_map_suffix (str): GT 的后缀,只读取特定后缀的 GT,默认值:'.png'
    ignore_index (int): 忽视掉的标签,默认值:255。比如我们可以把 GT padding 设为 255,这样这部分就不会被网络处理
    reduce_zero_label (bool): 忽视为 0 的标签,在实现中是将 0 变成 255,然后将其他所有标签值减一,默认值:False

添加索引

然后观察完之后得知 drive.pyhrf.pystare.py 这几个数据集是二分类数据集,因此数据的组织可以参考这几个。新建一个 mmseg/datasets/saliency.py

# mmseg/datasets/saliency.py
import os.path as osp

from .builder import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class SaliencyDataset(CustomDataset):
    """Saliency dataset.
    In segmentation map annotation for Saliency, 0 stands for background, which is
    included in 2 categories. ``reduce_zero_label`` is fixed to False. 
    """

    CLASSES = ('background', 'saliency')

    PALETTE = [[120, 120, 120], [6, 230, 230]]

    def __init__(self, **kwargs):
        super(SaliencyDataset, self).__init__(
            reduce_zero_label=False,
            **kwargs)
        assert osp.exists(self.img_dir)

然后将定义好的类添加到 mmseg/datasets/__init__.py

from .saliency.py import SaliencyDataset

__all__ = [
    ..., 'SaliencyDataset'
]

添加数据集配置

添加 configs/_base_/datasets/cod10k.py,内容依旧对着上面提到的二分类数据集改即可

# dataset settings
dataset_type = 'SaliencyDataset'
data_root = 'data/Saliency'  # 修改成你数据集所在路径
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
img_scale = (800, 800)  # 修改成你数据集合适的大小
crop_size = (384, 384)  # 修改成合适的大小
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),  # 这里不要 reduce_zero_label!!!
    dict(type='Resize', img_scale=img_scale, 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=img_scale,
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0],
        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,
    workers_per_gpu=4,
    train=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/training',  # 修改数据集路径,下同
        ann_dir='annotations/training',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline))

自定义训练配置

新建一个 configs/swin/upernet_swin_tiny_patch4_window7_384x384_40k_cod.py

_base_ = [  # 注意修改路径
    '../_base_/models/upernet_swin.py', '../_base_/datasets/cod10k.py',
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
]
model = dict(
    backbone=dict(
        embed_dim=96,
        depths=[2, 2, 6, 2],
        num_heads=[3, 6, 12, 24],
        window_size=7,
        ape=False,
        drop_path_rate=0.3,
        patch_norm=True,
        use_checkpoint=False
    ),
    decode_head=dict(
        in_channels=[96, 192, 384, 768],
        num_classes=2,  # 注意这里是两类
        # 由于正负样本非常不平衡,因此这里给正样本加了权重
        loss_decode=dict(class_weight=(1 ,8))
    ),
    auxiliary_head=dict(
        in_channels=384,
        num_classes=2,  # 注意这里是两类
        loss_decode=dict(class_weight=(1 ,8))
    ))

# AdamW optimizer, no weight decay for position embedding & layer norm in backbone
optimizer = dict(_delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.01,
                 paramwise_cfg=dict(custom_keys={'absolute_pos_embed': dict(decay_mult=0.),
                                                 'relative_position_bias_table': dict(decay_mult=0.),
                                                 'norm': dict(decay_mult=0.)}))

lr_config = dict(_delete_=True, policy='poly',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 power=1.0, min_lr=0.0, by_epoch=False)

# By default, models are trained on 8 GPUs with 2 images per GPU
# 修改了 batchsize 之后按比例修改 lr,bs 越大 lr 应该越大
data = dict(samples_per_gpu=5)

训练

tools/dist_train.sh configs/swin/upernet_swin_tiny_patch4_window7_384x384_40k_cod.py 2 --options model.pretrained=<PRETRAIN_MODEL>

其中 2 指的是 GPU 的数量,具体的还是看原项目的 README

一些报错

CUDA error: an illegal memory access was encountered

==mmsegmentation 中 gt 的范围是 [0, num_classes - 1]==,如果 gt 的值超过了这个范围就会出这个问题。由于增强等原因,可能前几个 ep 训练正常,到某个地方突然就报错。

KeyError: “‘XXX is not in the YYY registry’”

这个错误非常常见,我们只需要在 mmseg 下的相应文件里面去找即可。比如如果 YYY 是 backbone,那么就去 mmseg/models/backbones/__init__.py 里面去找,如果里面没有 XXX,那么要么是打错了,要么是不支持这个 backbone,也有可能没安装好。

第三种可能就输入 pip install -e . 命令重新安装即可。

Training loss is always 0.000

很有可能是类配置错了或者数据集配置有问题,对着上面的排查一下就好

RuntimeError: CUDA error: device-side assert triggered

检查一下是不是 num_classes 配置有问题

RuntimeError: DataLoader worker (pid 5657) is killed by signal: Killed

检查一下是不是爆内存了,检查一下数据集是不是有问题

跑 SETR 测试的时候出问题

vit 只支持固定大小的输入,正常的 test 流程出来的图片大小不一,会报错。

SETR 中在配置文件中设置 test_cfg = dict(mode='slide', crop_size=crop_size, stride=(320, 320)),不使用整个图像来预测,而是用固定大小的块来进行预测。

但是我们使用的数据集中有的图像较小,导致 crop 出来的还是大小不一,因此还需要在 configs/_base_/datasets/cod10k.py 里面修改 test_pipeline,具体来说就是在 Resize 之后再加上一个 Pad 操作:

transforms=[
    dict(type='Resize', ...),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    ...
]

其他问题

优先检查数据集是不是有问题,然后再检查配置文件是不是没对应上(尤其是类别数)。另外今天遇到一个问题是在 configs/_base_/datasets/cod10k.py 里面错误地设置了 reduce_zero_lable=True,导致训练出现了异常。