InternImage segmentation部分代码复现及训练自己的数据集（二）

一、训练自己的数据集

1、数据集准备

        准备好自己的训练数据集，因为只是做语义分割训练和测试，所以数据集应该包含图片数据和标注数据（主要是mask标注，用数字0~N-1标注N个类别，每一个类别用一个数字标注）。数据集的组织形式参考ADE20K数据集。或者有自己的组织形式也行，只要训练集和测试集分开，图片数据和标注数据分开就可以了。例如，下面为我自己的数据组织形式：

|--root_path
    |
    |--sign_semantic
         |
         |--train
         |    |
         |    |--Images
         |    |    |
         |    |    |--0000000023_0_0.png
         |    |    |--0000000023_0_1.png
         |    |    |--……
         |    |    
         |    |--Masks
         |    |    |
         |    |    |--0000000023_0_0.png
         |    |    |--0000000023_0_1.png
         |    |    |--……
         |    
         |--val
         |    |
         |    |--Images
         |    |    |
         |    |    |--0000001001_0_0.png
         |    |    |--0000001001_0_1.png
         |    |    |--……
         |    |    
         |    |--Masks
         |    |    |
         |    |    |--0000001001_0_0.png
         |    |    |--0000001001_0_1.png
         |    |    |--……
         |    |

2、代码修改

        这部分的代码主要涉及数据读取代码的修改和模型代码修改（改为你的数据集目标类别数目）。

（1）数据读取代码的修改

        参考上一篇文章：

InternImage segmentation部分代码复现及训练自己的数据集（一）_yeeeeeeeeeeees的博客-CSDN博客

        但是，上一篇文章中的方法一已经不适用了，在这里需要采取方法二的做法。参考文章：

(64条消息) mmsegmentation教程1:自定义数据集、config文件修改、训练教程_AESA相控阵的博客-CSDN博客

        首先下载安装mmsegmentation和mmcv-full扩展包，参考上一篇文章，这里不赘述。然后在环境路径下的 ./mmseg/datasets/ 下面对数据集进行初始定义，例如，我这里的环境路径是：

/home/lthpc/.conda/envs/internimage/lib/python3.7/site-packages/mmseg/datasets

        在 ./mmseg/datasets/ 文件夹下找到 stare.py 文件，复制一份，重命名为 my_dataset.py ，之所以选择 stare.py 文件，是因为我的语义分割也是做二分类。打开 my_dataset.py 文件进行编辑，下面带注释的部分为需要跟据自身的数据集进行修改的部分：

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .builder import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class SignDataset(CustomDataset):   #更改为你的数据集名称
    """STARE dataset.

    In segmentation map annotation for STARE, 0 stands for background, which is
    included in 2 categories. ``reduce_zero_label`` is fixed to False. The
    ``img_suffix`` is fixed to '.png' and ``seg_map_suffix`` is fixed to
    '.ah.png'.
    """

    CLASSES = ('background', 'sign')   #更改为你的数据集类别

    PALETTE = [[120, 120, 120], [6, 230, 230]]   #添加调色板，有多少个类别，就需要有多少组[R,G,B]像素值，用于在分割的时候进行填充

    def __init__(self, **kwargs):
        super(SignDataset, self).__init__(  #更改为你的数据集名称
            img_suffix='.png',   #更改为你的图片后缀
            seg_map_suffix='.png',  #更改为你的mask后缀
            reduce_zero_label=False,
            **kwargs)
        assert osp.exists(self.img_dir)

        然后，修改 ./mmseg/dataset/ 文件夹下的 __init__.py 文件，对数据集进行“注册（register）” 初始化，带注释的语句是需要添加或更改的：

# Copyright (c) OpenMMLab. All rights reserved.
from .ade import ADE20KDataset
from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
from .chase_db1 import ChaseDB1Dataset
from .cityscapes import CityscapesDataset
from .coco_stuff import COCOStuffDataset
from .custom import CustomDataset
from .dark_zurich import DarkZurichDataset
from .dataset_wrappers import (ConcatDataset, MultiImageMixDataset,
                               RepeatDataset)
from .drive import DRIVEDataset
from .face import FaceOccludedDataset
from .hrf import HRFDataset
from .imagenets import (ImageNetSDataset, LoadImageNetSAnnotations,
                        LoadImageNetSImageFromFile)
from .isaid import iSAIDDataset
from .isprs import ISPRSDataset
from .loveda import LoveDADataset
from .night_driving import NightDrivingDataset
from .pascal_context import PascalContextDataset, PascalContextDataset59
from .potsdam import PotsdamDataset
from .stare import STAREDataset
from .voc import PascalVOCDataset
from .my_dataset import SignDataset     #导入你的数据集对象

__all__ = [
    'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset',
    'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset',
    'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset',
    'PascalContextDataset59', 'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset',
    'STAREDataset', 'DarkZurichDataset', 'NightDrivingDataset',
    'COCOStuffDataset', 'LoveDADataset', 'MultiImageMixDataset',
    'iSAIDDataset', 'ISPRSDataset', 'PotsdamDataset', 'FaceOccludedDataset',
    'ImageNetSDataset', 'LoadImageNetSAnnotations',
    'LoadImageNetSImageFromFile', 'SignDataset'   #添加你的数据集进行注册
]

        然后，回到项目代码中，在  ./InternImage/segmentation/configs/_base_/datasets/ 文件夹下，找到 stare.py 文件，复制一份，重命名为 my_dataset.py，打开文件进行修改，下面带注释的语句是需要进行修改的：

# dataset settings
dataset_type = 'SignDataset'  #上一步定义的数据集名字
data_root = '/data/sign_semantic'   #数据集存储路径
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
img_scale = (605, 700)
crop_size = (224, 224)  #数据增强时裁剪的大小，这里跟据需要进行修改
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    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,  #batch_sign
    workers_per_gpu=4,  #nums gpu
    train=dict(
        type='RepeatDataset',
        times=40000,
        dataset=dict(
            type=dataset_type,
            data_root=data_root,
            img_dir='train/Images',   #训练图像路径
            ann_dir='train/Masks',  #训练mask路径
            pipeline=train_pipeline)),
    val=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='val/Images',   #验证集图像路径
        ann_dir='val/Masks',   #验证集mask路径
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='val/Images',  #测试集图像路径
        ann_dir='val/Masks',  #测试集mask路径
        pipeline=test_pipeline))

        至此，数据集读取底层文件的修改已经完成。

（2）模型修改

        这部分主要涉及两部分的修改，一是关联数据集读取文件，二是修改模型的输出。

        进入到 ./segmentation/configs/ade20k/ 文件夹下，找到 upernet_internimage_t_512_160k_ ade20k.py 文件，复制一份，重命名为 upernet_internimage_t_512_160k_sign.py，打开文件进行修改，修改如下面代码所示，需要修改的语句已经进行注释：

# --------------------------------------------------------
# InternImage
# Copyright (c) 2022 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
_base_ = [
    '../_base_/models/upernet_r50.py',   #这个是网络的backbone，使用单卡训练时要去backbone模型里将“SyncBN”改成“BN”
    '../_base_/datasets/my_dataset.py'  #换成自己定义的数据集
    '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_160k.py'
]
pretrained = 'https://huggingface.co/OpenGVLab/InternImage/resolve/main/internimage_t_1k_224.pth'
model = dict(
    backbone=dict(
        _delete_=True,
        type='InternImage',
        core_op='DCNv3',
        channels=64,
        depths=[4, 4, 18, 4],
        groups=[4, 8, 16, 32],
        mlp_ratio=4.,
        drop_path_rate=0.2,
        norm_layer='LN',
        layer_scale=1.0,
        offset_scale=1.0,
        post_norm=False,
        with_cp=False,
        out_indices=(0, 1, 2, 3),
        init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
    decode_head=dict(num_classes=2, in_channels=[64, 128, 256, 512]),   #通过改变这里的num_classes为你的数据集分类类别数（包括背景），完成配置
    auxiliary_head=dict(num_classes=2, in_channels=256),  #通过改变这里的num_classes为你的数据集分类类别数（包括背景），完成配置
    test_cfg=dict(mode='whole')
)
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(2048, 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='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
optimizer = dict(
    _delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.05,
    constructor='CustomLayerDecayOptimizerConstructor',
    paramwise_cfg=dict(num_layers=30, layer_decay_rate=1.0,
                       depths=[4, 4, 18, 4]))
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
data=dict(samples_per_gpu=2,
          val=dict(pipeline=test_pipeline),
          test=dict(pipeline=test_pipeline))
runner = dict(type='IterBasedRunner')
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=1)
evaluation = dict(interval=16000, metric='mIoU', save_best='mIoU')
# fp16 = dict(loss_scale=dict(init_scale=512))

        如果进行单卡训练，需要到网络的backbone部分修改BN层，参考：mmsegmentation教程1:自定义数据集、config文件修改、训练教程_AESA相控阵的博客-CSDN博客

        至此，需要修改的代码工作已经完成，后续网络的一些参数设置问题及代码注释，请关注更新。

3、训练测试

        完成上述的准备工作之后，终于可以开始训练你的代码了！

        终端回到 ./InternImage/segmentation/ 文件夹下，输入以下命令开始你的训练吧！

bash dist_train.sh configs/ade20k/upernet_internimage_t_512_160k_sign.py 2

        网络测试，请参考上一篇博客，或者 README.md 文件。