Vision Mamba代码笔记

原创

已于 2025-12-12 23:36:44 修改 · 1.2w 阅读

179 ·

CC 4.0 BY-SA版权

文章标签：

#笔记

于 2024-05-20 21:56:33 首次发布

1 论文回顾

基本思路

论文解读见：

《VideoMamba》论文笔记_video mamba-优快云博客

注意

Vision Mamba和VIT的输入和输出的shape的形状是相同的（VIT基于Transformer的Encoder设计，输入经过多层MHA和MLP计算，输入和输出的形状相同，Mamba的SSM架构就可以做到输入与输出token的个数以及每个token的维度相同，自然也可以做到整个输入和输出的形状相同，再者Vision Mamba的设计参照VIT的结构，自然也要注意输入与输出形状相同。两者的输入流经过各自对应的Encoder之后都具备了上下文信息，其效果相同，效率上基于Mamba的模型会更胜一筹。
正如1所说，Vision Mamba的设计参照VIT，这两个工作的流程是相同的，这里主要指的是图片打patch 再concat上class token再加上Position Embedding这个流程，两个模型唯一不同的地方就是Emcoder部分的不同，VIT使用的是Transformer的Encoder，Vim使用的是Mamba的Encoder，二者都是用于token间信息交互，上下文建模的

2 环境配置

按照官方readme.md配置，如果有问题照着下面这个链接改

vision mamba 运行训练记录，解决bimamba_type错误-优快云博客

值得说明的一点是，如果你之前在跑其他的mamba，环境拿过来是不能直接直接用的，因为标准的Mamba类是没有bimamba_type这个参数的，

所以，需要去Vim代码官网去找到mamba-1p1p1包，下载之后放自己项目里

事实上Vision Mamba重写了这个Mamba类，可以看到里边是由bimamba_type这个参数的（这其实也是Vision Mamba的主要贡献），执行如下代码

cp -rf mamba-1p1p1/mamba_ssm /home/liyhc/anaconda3/envs/mamba/lib/python3.10/site-packages
#后边是系统的mamba的安装路径，自己照着自己环境mamba的安装路径进行修改

3 代码笔记

3.1 代码链接

官方代码链接

Vim/vim/models_mamba.py at main · hustvl/Vim (github.com)

我手敲的带中文注释的链接

Johnny-Haytham/Vim: Vim with chinese notation (github.com)

3.2 Module

3.2.1 PatchEmbed

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224,patch_size=16,stride=16,in_channels=3,embed_dim=768,norm_layer=None,flatten=True):
        super(PatchEmbed, self).__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)#将img_size和patch_size化成元组的形式
        self.img_size = img_size
        self.patch_size = patch_size
        #一个patch形成一个grid（网格），这里记录网格的形状
        self.grid_size = ((img_size[0] - patch_size[0]) // stride + 1 , (img_size[1] - patch_size[1]) // stride + 1)
        self.num_patches = self.grid_size[0] * self.grid_size[1]#总共的patch个数
        self.flatten = flatten
        #打patch的操作，实际为卷积的操作(为了不重复卷积，步长的大小理论上因该等于卷积核的大小）
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()#nn.Identity的输入等于输出，通常作为占位层使用

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1],\
            f"Input img size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})"
        x = self.proj(x)#B,C,H,W——>B,embed_dim,grid_size,grid_size
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)#B,embed_dim,grid_size,grid_size——>B,embed_dim,grid_size*grid_size——>B,grid_size*grid_size,embed_dim
        x = self.norm(x)
        return x

3.2.2 Vim Encoder Block

class Block(nn.Module):
    def __init__(
            self, dim, mixer_cls,
            norm_cls = nn.LayerNorm,
            fused_add_norm=False,residual_in_fp32=False,drop_path=0.
    ):
        super(Block, self).__init__()
        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm

        self.mix