mmsegmentation教程2：如何修改loss函数、指定训练策略、修改评价指标、指定iterators进行val指标输出

简介

在mmseg教程1中对如何成功在mmseg中训练自己的数据集进行了讲解，那么能跑起来，就希望对其中loss函数、指定训练策略、修改评价指标、指定iterators进行val指标输出等进行自己的指定，下面进行具体讲解

具体修改方式

mm系列的核心是configs下面的配置文件，数据集设置与加载、训练策略、网络参数设置都在configs下面的配置文件里，可以简单理解为以前使用的开源框架中的main()函数，所有的参数都在这里指定，下面以configs/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K.py为例进行讲解

_base_ = [
    '../_base_/models/upernet_swin.py', 
    '../_base_/datasets/egc.py',
    '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]
model = dict(
    pretrained='pretrain/swin_tiny_patch4_window7_224.pth',
    backbone=dict(
        embed_dims=96,
        depths=[2, 2, 6, 2],
        num_heads=[3, 6, 12, 24],
        window_size=7,
        use_abs_pos_embed=False,
        drop_path_rate=0.3,
        patch_norm=True),
    decode_head=dict(in_channels=[96, 192, 384, 768], num_classes=150),
    auxiliary_head=dict(in_channels=384, num_classes=150))

# 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)

evaluation = dict(interval=100, metric='mIoU', pre_eval=True)

# By default, models are trained on 8 GPUs with 2 images per GPU
data = dict(
    samples_per_gpu=16,
    workers_per_gpu=4,
)

下面进行拆开讲解：

_base_讲解

_base_ = [
    '../_base_/models/upernet_swin.py', 
    '../_base_/datasets/egc.py',
    '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

upernet_swin.py：表示网络模型定义
egc.py：表示你要使用的数据集
default_runtime.py：表示运行策略
schedule_80k.py：表示训练策略

整个过程可以理解为导包

model部分讲解

model = dict(
    pretrained=None,
    backbone=dict(
        embed_dims=96,
        depths=[2, 2, 6, 2],
        num_heads=[3, 6, 12, 24],
        window_size=7,
        use_abs_pos_embed=False,
        drop_path_rate=0.3,
        patch_norm=True),
    decode_head=dict(in_channels=[96, 192, 384, 768], num_classes=150),
    auxiliary_head=dict(in_channels=384, num_classes=150))

这一部分主要是针对configs/_ base _/models/upernet_swin.py

# model settings
norm_cfg = dict(type='BN', requires_grad=True)
backbone_norm_cfg = dict(type='LN', requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained=None,
    backbone=dict(
        type='SwinTransformer',
        pretrain_img_size=224,
        embed_dims=96,
        patch_size=4,
        window_size=7,
        mlp_ratio=4,
        depths=[2, 2, 6, 2],
        num_heads=[3, 6, 12, 24],
        strides=(4, 2, 2, 2),
        out_indices=(0, 1, 2, 3),
        qkv_bias=True,
        qk_scale=None,
        patch_norm=True,
        drop_rate=0.,
        attn_drop_rate=0.,
        drop_path_rate=0.3,
        use_abs_pos_embed=False,
        act_cfg=dict(type='GELU'),
        norm_cfg=backbone_norm_cfg),
    decode_head=dict(
        type='UPerHead',
        in_channels=[96, 192, 384, 768],
        in_index=[0, 1, 2, 3],
        pool_scales=(1, 2, 3, 6),
        channels=512,
        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)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=384,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

进行复写，其实很好理解，upernet_swin.py只是模型的初始定义，你要是使用模型就需要修改一些东西，比如你的分类数num_classes有可能不一样、你想换一个loss进行训练，就需要对upernet_swin.py进行复写，这里复写一定要注意格式，一定要保证复写中各个部件的层级关系一定要与upernet_swin.py的一样，比如backbone是在model的dict之下的，decode_head是与backbone平级的。
通常情况下复写只需要对三个地方进行复写：

pretrained：是否使用预训练数据集
decode_head下面的num_classes
decode_head下面的loss_decode下面的type中的损失函数

对num_classes进行修改

这个很简单，你把decode_head下面的num_classes改为你需要的分类类别数量就行。

对loss函数进行修改

mmseg中的loss函数定义在mmseg/models/losses/_ int _.py下面

# Copyright (c) OpenMMLab. All rights reserved.
from .accuracy import Accuracy, accuracy
from .cross_entropy_loss import (CrossEntropyLoss, binary_cross_entropy,
                                 cross_entropy, mask_cross_entropy)
from .dice_loss import DiceLoss
from .focal_loss import FocalLoss
from .lovasz_loss import LovaszLoss
from .utils import reduce_loss, weight_reduce_loss, weighted_loss

__all__ = [
    'accuracy', 'Accuracy', 'cross_entropy', 'binary_cross_entropy',
    'mask_cross_entropy', 'CrossEntropyLoss', 'reduce_loss',
    'weight_reduce_loss', 'weighted_loss', 'LovaszLoss', 'DiceLoss',
    'FocalLoss'
]

基本上常用的都有了，可以自行选择，比如在upernet_swin.py中的loss函数默认是

loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),  

复写的时候你想改成DiceLoss，那么复写的时候加上

loss_decode=dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0)),  

就OK了。

对训练策略部分进行讲解

训练策略对应

../_base_/schedules/schedule_80k.py

对应的是

configs/_base_/schedules/schedule_80k.py

长这样

# optimizer
#torch.optim下的优化器，可选[SGD, Adam, AdamW, RMSprop, Adamax]
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
optimizer_config = dict()
# learning policy
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
# runtime settings
runner = dict(type='IterBasedRunner', max_iters=80000)
checkpoint_config = dict(by_epoch=False, interval=8000)
evaluation = dict(interval=8000, metric='mIoU', pre_eval=True)

optimizer ：表示优化器，可以使用torch.optim下的优化器，通常使用[SGD, Adam, AdamW, RMSprop, Adamax]

lr_config： 学习率策略，默认为poly

runner： 训练周期，IterBasedRunner表示基于iterators进行训练，最大iterators=80K

checkpoint_config：by_epoch表示是否按照epoch的方式进行记录，interval表示每多少个周期保存一次输出结果

evaluation： 表示验证测量，interval表示多少个interval进行一次验证并输出验证结果，默认每8K个iterators进行一次验证；metric表示你使用的验证指标，默认为mIoU。

同样，这里面的东西也都是可以改，比如上面就改成了这样

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)

evaluation = dict(interval=100, metric='mIoU', pre_eval=True)

把优化器从SGD换成了AdamW

为了更加方便的的8K个iterators才进行一次验证太久了，换成每100个进行一次验证

下面重点介绍评价指标的更改

更改评估指标

mmseg默认的评估指标是mIoU，评价指标在
mmseg/core/evaluation/metrics.py下面的

def total_area_to_metrics(total_area_intersect,
                          total_area_union,
                          total_area_pred_label,
                          total_area_label,
                          metrics=['mIoU'],
                          nan_to_num=None,
                          beta=1):
    
    if isinstance(metrics, str):
        metrics = [metrics]
    allowed_metrics = ['mIoU', 'mDice', 'mFscore']
    if not set(metrics).issubset(set(allowed_metrics)):
        raise KeyError('metrics {} is not supported'.format(metrics))

    all_acc = total_area_intersect.sum() / total_area_label.sum()

	#无论你选谁都会输出，aAcc其实就是mPA（平均像素准确度）
    ret_metrics = OrderedDict({'aAcc': all_acc})
    
	'''重点看这下面！！！！！！'''
    for metric in metrics:
    	#选取mIoU为指标
        if metric == 'mIoU':
            iou = total_area_intersect / total_area_union
            acc = total_area_intersect / total_area_label
            ret_metrics['IoU'] = iou
            ret_metrics['Acc'] = acc
        elif metric == 'mDice':
            dice = 2 * total_area_intersect / (
                total_area_pred_label + total_area_label)
            acc = total_area_intersect / total_area_label
            ret_metrics['Dice'] = dice
            ret_metrics['Acc'] = acc
        elif metric == 'mFscore':
            precision = total_area_intersect / total_area_pred_label
            recall = total_area_intersect / total_area_label
            f_value = torch.tensor(
                [f_score(x[0], x[1], beta) for x in zip(precision, recall)])
            ret_metrics['Fscore'] = f_value
            ret_metrics['Precision'] = precision
            ret_metrics['Recall'] = recall

    ret_metrics = {
        metric: value.numpy()
        for metric, value in ret_metrics.items()
    }
    if nan_to_num is not None:
        ret_metrics = OrderedDict({
            metric: np.nan_to_num(metric_value, nan=nan_to_num)
            for metric, metric_value in ret_metrics.items()
        })
    return ret_metrics

可以看到，选择mIoU为评价指标的同时会输出各个类别的IoU，同样会给你输出acc，其实这里的acc就是PA（像素准确度），加上无论你选谁都会输出的aAcc（其实就是mPA（平均像素准确度））。

所以你得到指标一共有
mIoU，IoU，acc, aAcc

你要是想换mDice为指标，那简单，把mIoU直接换成mDice就行，如果想成年人不做选择，两个都要，那也简单，直接在mIoU下面加入mDice的计算就行，比如下面这样：

'''重点看这下面！！！！！！'''
    for metric in metrics:
    	#选取mIoU为指标
        if metric == 'mIoU':
            iou = total_area_intersect / total_area_union
            #新加的dice
            dice = 2 * total_area_intersect / (
                total_area_pred_label + total_area_label)
            acc = total_area_intersect / total_area_label
            ret_metrics['IoU'] = iou
            #新加的dice
            ret_metrics['Dice'] = dice
            ret_metrics['Acc'] = acc
        elif metric == 'mDice':
            dice = 2 * total_area_intersect / (
                total_area_pred_label + total_area_label)
            acc = total_area_intersect / total_area_label
            ret_metrics['Dice'] = dice
            ret_metrics['Acc'] = acc

这样你会得到：
mIoU，IoU，mDice, Dice, acc, aAcc

对数据集加载部分进行讲解

数据集加载对应的是

../_base_/datasets/egc.py

这里对应的是

configs/_base_/datasets/egc.py

和model同样，你可以对egc.py(这是我的数据集定义的名字)任意进行复写，只要保证层级关系对就行。
这里通常只需要复写一个东西

data = dict(
    samples_per_gpu=16,
    workers_per_gpu=4,
)

samples_per_gpu：表示每张GPU的样本数，可以简单理解为batch_size
workers_per_gpu： 表示多线程数

这两个东西共同组成batch_size，运算关系大概是

epoch = (iterators * batch_size) / (图片总数)

这里需要关注的地方在于，如果显存不足你就需要降低samples_per_gpu，但是由于mmseg不是使用epoch，而是使用iterators ，所以你如果把samples_per_gpu降低了2倍，相应的需要在iterators *2，比如samples_per_gpu从16变成8，你就需要

data = dict(
    samples_per_gpu=8,
    workers_per_gpu=2,
)
index = 16 / 8
# runtime settings
runner = dict(type='IterBasedRunner', max_iters=80000 * index)
checkpoint_config = dict(by_epoch=False, interval=8000 * index)
evaluation = dict(interval=100, metric='mIoU', pre_eval=True)