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

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&#ister_module()
class SignDataset(CustomDataset):   #更改为你的数据集名称"""STARE dataset.In segmentation map annotation for STARE, 0 stands for background, which isincluded 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)ists(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_signworkers_per_gpu=4,  #nums gputrain=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 = '.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 文件。