mmseg代码从头修改-1

mmseg代码从头修改-1

问题来源

​ 由于我的这段科研工作主要围绕上下采样同时建模（协同上下采样）来展开的，我在小模型上跑出了很好的效果，想在大模型大数据集上进一步验证他的有效性，我选择了 mask2former 模型作为分割任务的 baseline，我尝试从 mask2former 官方 repo 下载代码后在他上面修改，但是这个代码封装极其深而且工程化，这个仓库是基于一个 detectron2 这个算法框架库开发的，首先解析这个算法框架库就极其耗时，而且由于耦合极其深也不好修改。这时候我发现 mmseg 这个框架库内有 mask2former 模型而且这个框架更好理解一些（但也是耦合太深封装过度），我就决定在 mmseg 这个算法框架下修改代码加入协同算子

​ 这里引出了第二个问题，如果我只是单纯地做一个上采样算子就非常好说了，但是我将上下采样算子同时建模了，而 mmseg 这个框架库实例化模型的时候是将 backbone(ResNet50) 和 head(FPNHead) 分别生成的（解耦了），而我的协同算子既不能单独属于 backbone 也不能单独属于 head，这就意味着我的 Sampler 需要和 backbone 和 head 平级来实例化

修改方案设计

​ 我们先看在 mmseg 内部 mask2former 的 Config 配置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
model = dict(
    type='EncoderDecoder',
    data_preprocessor=data_preprocessor,
    backbone=dict(
        type='ResNet',
        ...),
    


# 下面一长串全是 decode_head
    decode_head=dict(
        type='Mask2FormerHead',
        in_channels=[256, 512, 1024, 2048],
        strides=[4, 8, 16, 32],
        feat_channels=256,
        out_channels=256,
        num_classes=num_classes,
        num_queries=100,
        num_transformer_feat_level=3,
        align_corners=False,
        
        # 第一 decoder，要改为FPN
        pixel_decoder=dict(
            type='mmdet.MSDeformAttnPixelDecoder',
            ...),
       
        enforce_decoder_input_project=False, 
        positional_encoding=dict(  # SinePositionalEncoding
            num_feats=128, normalize=True),
        
        
        # 第二个 decoder，保留不变
        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
            ...),
        
        loss_cls=dict(...),
        loss_mask=dict(...),
        loss_dice=dict(...),
        train_cfg=dict(...))



    train_cfg=dict(),
    test_cfg=dict(mode='whole')
)

对应的论文中网络结构图如下：

network

​

​ 可以注意到，上面的论文把 ResNet 称为 backbone 而不是 ResNet+Pixel Decoder，对应 config 的层级关系是：backbone - Pixel Decoder - Transformer Decoder，而我们需要将 backbone 与 Pixel Decoder 同时考虑进来（因为要加入上下采样管理），因此我们需要的层级关系是 (backbone - Pixel Decoder) - Transformer Decoder，即把 backbone 与 pixel decoder 合起来的一级和 Transformer Decoder 同级

​ 那么最简单的设计就是把Backbone和Pixel Decoder封装为一个新的Backbone，Transfomer De coder就当做唯一的Decoder就行了！

  • 为什么不把(backbone - Pixel Decoder)视为一个单独的特殊模块而是 backbone？这是因为如过把他视为一个新的模块，我们要去 registry 文件（mmengine）中再加上实例化这个新的模块的代码，如果把他视为一个 backbone 就只需要新写一个 backbone 再调用 backbone 的实例化代码即可

即不需要这样麻烦的操作：

1
CUSTOM_MODULES = Registry('custom modules')