VSSM VMamba实现

文章目录

VSSM

Mamba实现可以参照之前的
mamba_minimal系列
论文地址：
VMamba
论文阅读：
VMamba:视觉状态空间模型
代码地址：
https://github.com/MzeroMiko/VMamba.git
SS2D实现

以分类任务用到的VMamba为例。

维度变换

操作的具体参数定义见初始化

阶段	维度
输入x	$[B, C, H, W]$
embed	$B, H/4, W/4, C_1 ]$
阶段1	$B, H/4, W/4, C_1 ]$
阶段2	$B, H/8, W/8, C_2 ]$
阶段3	$B, H/16, W/16, C_3 ]$
阶段4	$B, H/32, W/32, C_4 ]$
分类器	$[B, 1000]$

在这里插入图片描述

初始化

参数	定义	说明
in_chans	3	输入图像的通道数
depths	[2, 2, 9, 2]	定义每层的VSS Block数
dims	[96, 192, 384, 768]	定义每层的输出通道数
downsample_version	v2	下采样操作的版本
patchembed_version	v1	图像嵌入
mlp_ratio	4.0	定义mlp隐藏维度缩放
ssm_d_state	16	ssm隐状态的维度
ssm_ratio	2.0	d_inner = d_state * ssm_ratio
ssm_init	v0	ssm初始化版本
forward_type	v2	ssm前向版本

模型参数初始化

大部分参数即SS2D，VSS块中的参数由定义的ssm初始化版本初始化，剩下的线性层和归一化层参数由下面的函数初始化。

    def _init_weights(self, m: nn.Module):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
1
2
3
4
5
6
7
8

模型搭建

def_make_layer

构建VSSM的4个阶段，即4层，VSSBlock本身并不改变输入的尺寸，因此需要下采样模块将输出维度变换为下一阶段的输入维度

def _make_layer(
        dim=96, 
        drop_path=[0.1, 0.1], 
        use_checkpoint=False, 
        norm_layer=nn.LayerNorm,
        downsample=nn.Identity(),
        # ===========================
        ssm_d_state=16,
        ssm_ratio=2.0,
        ssm_dt_rank="auto",       
        ssm_act_layer=nn.SiLU,
        ssm_conv=3,
        ssm_conv_bias=True,
        ssm_drop_rate=0.0, 
        ssm_init="v0",
        forward_type="v2",
        # ===========================
        mlp_ratio=4.0,
        mlp_act_layer=nn.GELU,
        mlp_drop_rate=0.0,
        **kwargs,
    ):
        depth = len(drop_path)
        blocks = []
        for d in range(depth):
            blocks.append(VSSBlock(
                hidden_dim=dim, 
                drop_path=drop_path[d],
                norm_layer=norm_layer,
                ssm_d_state=ssm_d_state,
                ssm_ratio=ssm_ratio,
                ssm_dt_rank=ssm_dt_rank,
                ssm_act_layer=ssm_act_layer,
                ssm_conv=ssm_conv,
                ssm_conv_bias=ssm_conv_bias,
                ssm_drop_rate=ssm_drop_rate,
                ssm_init=ssm_init,
                forward_type=forward_type,
                mlp_ratio=mlp_ratio,
                mlp_act_layer=mlp_act_layer,
                mlp_drop_rate=mlp_drop_rate,
                use_checkpoint=use_checkpoint,
            ))
        
        return nn.Sequential(OrderedDict(
            blocks=nn.Sequential(*blocks,),
            downsample=downsample,
        ))
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
46
47
48

def _make_downsample

默认下采样版本v2

下采样模块，通过2D卷积之后，长宽变为原来的一半，通道数不变

    def _make_downsample(dim=96, out_dim=192, norm_layer=nn.LayerNorm):
        return nn.Sequential(
            Permute(0, 3, 1, 2),
            nn.Conv2d(dim, out_dim, kernel_size=2, stride=2),
            Permute(0, 2, 3, 1),
            norm_layer(out_dim),
        )
1
2
3
4
5
6
7

patch embed

默认嵌入版本v1，对输入图像进行embed

输入x维度 $[B, 3, H, W]$ ，嵌入后通道维变为96， $\frac{H}{patch\_size}$ ， $\frac{W}{patch\_size}$ ， $\frac{H}{4}, \frac{W}{4}]$

 def _make_patch_embed(in_chans=3, embed_dim=96, patch_size=4, patch_norm=True, norm_layer=nn.LayerNorm):
        return nn.Sequential(
            nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True),
            Permute(0, 2, 3, 1),
            (norm_layer(embed_dim) if patch_norm else nn.Identity()), 
        )
1
2
3
4
5
6

第一至四阶段

这几个阶段的差别在于每一层的VSSBlock数不同，由depths定义分别为 [2, 2, 9, 2]，输出维度由dims定义分别为[96, 192, 384, 768]。其组成元素除一阶段外，均在VSSBlock前包含下采样模块以变换维度。

具体介绍见VSSBlock

分类器

池化后长宽变为1，则变量尺寸变为 $[B, C, 1, 1]$ ，展平后变为 $[B, C]$ 最后线性投影到类别维度1000

$[B, 1000]$

self.classifier = nn.Sequential(OrderedDict(
            norm=norm_layer(self.num_features), # B,H,W,C
            permute=Permute(0, 3, 1, 2),
            avgpool=nn.AdaptiveAvgPool2d(1),
            flatten=nn.Flatten(1),
            head=nn.Linear(self.num_features, num_classes),
        ))
1
2
3
4
5
6
7

VSSBlock

对于ssm分支来说，其输入输出维度不变为(B, H, W, d_model) ，对于mlp分支来说中间的隐藏维度根据mlp_ratio参数定义会有所增加，但是最后又会映射为原来的维度，因此整体上并不改变输入的维度。

def init

主要分为两个分支ssm分支和mlp分支

ssm分支

主要组成部分是SS2D块
SS2D实现

if self.ssm_branch:
            self.norm = norm_layer(hidden_dim)
            self.op = _SS2D(
                d_model=hidden_dim, 
                d_state=ssm_d_state, 
                ssm_ratio=ssm_ratio,
                dt_rank=ssm_dt_rank,
                act_layer=ssm_act_layer,
                # ==========================
                d_conv=ssm_conv,
                conv_bias=ssm_conv_bias,
                # ==========================
                dropout=ssm_drop_rate,
                # =========================
                initialize=ssm_init,
                forward_type=forward_type,
            )      
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

图中的SS2D和SS2D类的定义有偏差，简单来说是是包含SS2D块加一个残差连接，图中所示SS2D应表示状态空间模型SSM部分，即VSS块相比SS2D块只增加了残差连接和入口的归一化。如果定义了MLP分支，VSS块的输出还会经过一个残差连接的两层MLP