企业官网建设流程全解析

医学图像分割这个方向，做过的都知道，精度和显存基本是一对冤家。UNet 系列靠卷积堆局部感受野，Transformer 系列靠自注意力拿全局建模能力，但自注意力的计算量随分辨率平方增长，一张 512×512 的 CT 切片还好，换成 1024×1024 的病理全片或者 3D 的 MRI 体数据，显存直接告急。Mamba 这条路线之所以在视觉圈火起来，核心就一句话：它用状态空间模型（SSM）做序列建模，复杂度是线性的，同时还能保留类似注意力的长程依赖能力。VMamba 就是把这套东西搬进视觉主干的代表作，而它最关键的视觉适配改造，就是 SS2D 四向扫描。这篇就围绕医学图像分割这个具体场景，把 VMamba 的 SS2D 机制、为什么这么设计、怎么落地、踩过哪些坑，一次讲透。

1. 医学图像分割为什么盯上了状态空间模型

1.1 卷积与自注意力在分割任务里的两难

先把这个需求背景说清楚，不然后面讲 SS2D 会显得像空中楼阁。医学图像分割的典型输入是 CT、MRI、超声、病理切片，输出是逐像素的类别掩码。它和自然图像分割最大的区别在于：目标边界往往模糊、对比度低、器官形状个体差异大，而且经常需要 3D 上下文。比如肝脏肿瘤分割，肿瘤和周围正常肝组织的灰度差可能只有十几个 HU，单看局部窗口根本分不出来，必须结合较大范围的解剖结构关系。

卷积网络的归纳偏置是局部性，堆深了感受野确实会变大，但它是各向同性的、固定的，对"这个像素属于肿瘤还是血管"这种需要长程语义判断的问题，卷积的感受野增长效率很低。Transformer 的全局自注意力理论上一步到位，但代价是 O(N²)。一张 512×512 的切片展平成 262144 个 token，自注意力矩阵就是 262144²，这个量级没有任何单卡能扛住，实际都得靠窗口注意力或者下采样来妥协，而一旦加窗口，全局性又打了折扣。

所以这个领域一直在等一个"既要长程建模、又要线性复杂度"的架构。状态空间模型正好卡在这个位置上。

1.2 状态空间模型的一句话直觉

状态空间模型（State Space Model, SSM）本质上是描述一个连续系统如何随时间演化的微分方程。用最朴素的话讲：它维护一个内部状态 h(t)，每来一个新的输入 x(t)，状态就按某种规则更新，同时吐出一个输出 y(t)。连续形式写作 h'(t) = A·h(t) + B·x(t)，y(t) = C·h(t)。离散化之后变成递推形式 h_t = Ā·h_{t-1} + B̄·x_t，y_t = C·h_t。

这个递推结构有个极其重要的性质：当前状态 h_t 压缩了此前所有输入的信息，所以它天然具备"记忆"，而且每一步的计算量是固定的，跟序列长度无关。整个序列扫一遍，总复杂度就是 O(N)。这就是线性复杂度的来源。对比 RNN，它同样是递推，但 SSM 的 A、B、C 参数是通过 HiPPO 等理论构造的，能让状态稳定地记住长程信息，不像普通 RNN 那样梯度爆炸或遗忘。

Mamba 的贡献在于把这个 A、B、C 做成了输入相关的（selective scan），也就是参数随当前 token 动态变化，这让模型能"选择性地"记住或忽略信息，表达能力大幅提升。但 selective scan 有个副作用：它不再是时不变的，没法用普通卷积来并行加速，必须用专门的并行扫描算法（parallel scan）在 GPU 上高效实现。这是 Mamba 工程落地的核心难点，也是后面讲 SS2D 实现时要重点说的。

1.3 从一维序列到二维图像：方向缺失问题

Mamba 原生是为一维序列设计的，语言模型里 token 有明确的前后顺序，扫一遍就够了。但图像是二维的，一个像素的邻居有上下左右四个方向，你把它按行优先展平成一维，那纵向的相邻关系就被打散到相隔一整行的位置去了。如果只做单向扫描，模型对"下方"信息的感知会非常弱，这对分割任务是致命的——器官的上下边界、肿瘤的浸润方向，都依赖二维空间关系。

这就是 VMamba 要解决的核心问题：怎么把一维的 SSM 扫描，改造成能覆盖二维空间四个方向的形式。答案就是 SS2D（2D Selective Scan），四向扫描。下面拆开讲。

2. SS2D 四向扫描到底在扫什么

2.1 四条扫描路径的展开方式

SS2D 的做法说起来不复杂，但细节很讲究。它把一张 H×W 的特征图，沿着四个方向分别展平成四条序列：

• 方向一：从左上到右下，逐行扫描（行优先，正向）
• 方向二：从右下到左上，逐行反向扫描
• 方向三：从左上到右下，逐列扫描（列优先，正向）
• 方向四：从右下到左上，逐列反向扫描

每条序列独立送进 S6（selective scan）模块做状态空间递推，得到四个输出序列，然后把它们重新 reshape 回 H×W 的特征图，最后相加或者拼接融合。这样做的效果是：任意一个像素，在四条路径里至少有一条路径上，它的上下左右邻居是紧挨着它的。比如列扫描保证了纵向相邻像素在序列里是连续的，行扫描保证了横向相邻像素连续，正反两个方向则保证了信息既能从前往后传，也能从后往前传。

我一开始以为四向扫描是"锦上添花"，实测下来不是。在肝脏分割任务上，把四向砍成单向，Dice 直接掉了三到四个点，边界处的锯齿明显变多。原因就是单向扫描丢失了垂直方向的连续性，模型只能靠卷积层去补，补不回来。

2.2 为什么是四个方向而不是八个

有人会问，二维有八个方向（加上两条对角线），为什么 VMamba 只做四个？这里有个工程上的权衡。对角线方向的相邻像素在展平后本来就不是严格连续的，加了对角线扫描，序列里相邻的两个 token 在空间上可能是斜对角，状态递推的"局部连续性"假设就被破坏了，收益不明显反而增加计算。四个正交方向已经能覆盖所有水平垂直的邻接关系，对角线关系可以通过两层扫描的叠加间接建模。实测加对角线版本，参数量和显存都涨了，精度没有稳定提升，所以主流实现都停在四向。

2.3 四向融合的方式选择

四个方向的输出怎么合并，也有讲究。常见做法有两种：逐元素相加，或者通道拼接后接一个线性投影。相加的好处是参数量零增加、计算量小，坏处是四个方向的信息被强制对齐到同一通道语义，可能互相干扰。拼接再投影的好处是让网络自己学怎么加权融合，表达更灵活，代价是多一个 4C→C 的线性层。

VMamba 官方实现用的是相加。我在自己的分割项目里试过拼接版本，在数据量充足（上万张切片）时拼接略好一点点，但在小数据集（几百张）上相加更稳，因为拼接引入的额外参数容易过拟合。所以如果你做的是小样本医学分割，老老实实用相加。

3. 把 SS2D 塞进分割网络的完整搭建思路

3.1 整体骨架：VMamba 主干加解码器

VMamba 本身是个分类主干，做分割要自己接解码器。我的做法是拿 VMamba 的四阶段特征（stride 4/8/16/32），像 UNet 那样做 skip connection，解码器用简单的卷积上采样加特征拼接。这里有个关键点：VMamba 的 stage 之间是下采样加通道翻倍，和 Swin 类似，所以解码器可以直接复用 Swin-UNet 那套结构，改动量很小。

具体通道配置我一般设成 [96, 192, 384, 768]，和官方 tiny 版本对齐。医学图像分辨率通常比自然图像小（512 或 256），所以第一个 stage 的 stride 我有时会从 4 改成 2，保留更多细节，代价是显存涨一截。这个取舍要看你的 GPU 和任务对边界精度的要求。

3.2 VSS Block 的内部结构

SS2D 是包在 VSS Block 里的。一个 VSS Block 的数据流是这样的：输入先过 LayerNorm，然后分两路，一路进 SS2D 做空间建模，另一路过一个线性层做门控（类似 GELU 激活的支路），两路逐元素相乘，最后残差加回输入。这个门控结构借鉴了 Gated Linear Unit 的思路，让 SS2D 的输出能被动态调制。

这里有个容易忽略的细节：SS2D 内部其实还包含一个深度可分离卷积（DWConv），放在 selective scan 之前。它的作用是在扫描前给局部邻域做一次信息混合，弥补 SSM 递推本身对局部细节建模偏弱的问题。这个 DWConv 的卷积核一般是 3×3，别小看它，去掉之后分割边界会明显变糊。

3.3 关键代码路径拆解

下面这段是 SS2D 前向的核心逻辑，我按自己的理解加了注释，帮你理清数据形状的变化：

import torch import torch.nn as nn class SS2D(nn.Module): def __init__(self, d_model, d_state=16, ssm_ratio=2.0): super().__init__() self.d_model = d_model self.d_state = d_state self.d_inner = int(ssm_ratio * d_model) # 输入投影：把通道分成 x 和 z 两路，z 用于门控 self.in_proj = nn.Linear(d_model, self.d_inner * 2) # 深度可分离卷积，扫描前做局部混合 self.conv2d = nn.Conv2d( self.d_inner, self.d_inner, kernel_size=3, padding=1, groups=self.d_inner ) # SSM 参数：dt、A、B、C 的投影 self.x_proj = nn.Linear(self.d_inner, d_state * 2 + 1) self.dt_proj = nn.Linear(d_state, self.d_inner) # 输出投影 self.out_proj = nn.Linear(self.d_inner, d_model) def forward(self, x): # x: (B, H, W, C) B, H, W, C = x.shape xz = self.in_proj(x) # (B, H, W, 2*d_inner) x, z = xz.chunk(2, dim=-1) # 各 (B, H, W, d_inner) x = x.permute(0, 3, 1, 2) # (B, d_inner, H, W) x = self.conv2d(x) # 局部混合 x = x.permute(0, 2, 3, 1) # 回到 (B, H, W, d_inner) # 四向扫描：这里用四次不同 reshape 模拟四个方向 y = self.four_direction_scan(x) y = y * torch.nn.functional.silu(z) # 门控 return self.out_proj(y)

four_direction_scan是核心，它内部对每个方向做 reshape、selective scan、再 reshape 回来。真正的 selective scan 一般调用 CUDA 实现（比如 mamba_ssm 库里的selective_scan_fn），纯 PyTorch 版本能跑但慢很多，只适合调试。

4. 医学分割场景下的实操配置与调参

4.1 输入尺寸与 patch 策略

医学图像分割很少直接喂全图，尤其是病理和 3D 数据。我的常规做法是：2D 切片用 512×512 随机裁剪，3D 体数据用 128×128×128 的 patch。VMamba 对输入尺寸没有硬性要求，但 SS2D 的显存占用和 H×W 成正比（线性复杂度），所以尺寸放大一倍，显存大概翻一倍，比 Transformer 的平方增长友好太多。这也是它在高分辨率医学图像上的最大优势。

有个坑要提醒：SS2D 的四向扫描在 H 和 W 差距很大时（比如 1024×64 的条带），列扫描的序列会特别长，行扫描特别短，负载不均衡。如果必须用非方形输入，建议先 pad 成接近方形再裁回来。

4.2 学习率与优化器

VMamba 主干对学习率比较敏感。我用 AdamW，主干学习率设 1e-4，解码器设 1e-3（差一个数量级），weight decay 0.05。如果主干也用 1e-3，训练前几个 epoch loss 会震荡甚至发散，因为 SSM 的递推结构对参数扰动比卷积敏感。warmup 建议 5 个 epoch，cosine 衰减到底。

另外 SSM 里的 dt 参数初始化很关键，官方实现里 dt 的 bias 初始化决定了初始的"记忆衰减速度"。如果发现模型学不到长程依赖（表现为大器官分割还行、细长结构如血管一塌糊涂），可以试着调大 dt 的初始 bias，让状态衰减慢一点，记忆更长。

4.3 数据增强的取舍

医学图像增强和自然图像不一样，不能随便翻转旋转，因为器官有解剖方位。我的经验是：水平翻转可以用（左右对称器官除外），垂直翻转慎用，旋转限制在 ±15 度以内，弹性形变对软组织分割很有帮助但幅度要小。颜色抖动对 CT/MRI 意义不大，但对病理 H&E 染色切片很有用，可以做色调和饱和度的轻微扰动。

这里有个和 SS2D 相关的注意点：四向扫描对方向是敏感的，如果你做了大量翻转增强，模型需要额外学习方向不变性，收敛会慢一些。实测在翻转增强很强的情况下，把四向扫描改成"两向加共享参数"能稍微缓解，但会损失一点精度，看你怎么权衡。

5. 踩过的坑与排查链路

5.1 显存没降反升的诡异现象

第一次跑 VMamba 分割，我满心以为线性复杂度能省显存，结果 512×512 输入下显存比 Swin 还高。排查了一圈发现两个原因：一是 SS2D 的中间激活（四条扫描序列）都要保留用于反向传播，激活值总量是输入的四倍；二是 selective scan 的 CUDA 实现内部有额外的 workspace 开销。

解决办法有三个：用 gradient checkpointing 把 SS2D 的中间激活换时间；把 ssm_ratio 从 2.0 降到 1.0（d_inner 减半）；减小 batch size 配合梯度累积。我最后用的是 checkpointing 加 batch 8 加累积 4 步，等效 batch 32，显存压到了 24G 以内。

5.2 训练 loss 正常但 Dice 上不去

这个坑更隐蔽。loss 曲线很漂亮，验证 loss 也在降，但 Dice 卡在 0.7 上不去。后来把预测结果可视化才发现，模型把大部分像素都预测成了背景，因为背景占比太高（医学图像典型的前景占比不到 5%），交叉熵被背景主导了。

换成 Dice loss 加交叉熵的混合损失后好转，但还没到最优。真正解决问题的是加了边界加权：在器官边界处给更高的 loss 权重。因为 SS2D 虽然能建模长程依赖，但对精细边界的定位还是偏弱，需要显式引导。这个技巧在医学分割里几乎是标配，但很多人一开始会忽略。

5.3 四向扫描的 CUDA 编译问题

mamba_ssm 的 selective scan 需要编译 CUDA 扩展，环境不对会报一堆错。最常见的两个：一是 CUDA 版本和 PyTorch 版本不匹配，二是 gcc 版本太新导致编译失败。我的建议是严格按官方要求的版本组合来，PyTorch 2.1 配 CUDA 11.8 是比较稳的组合。如果实在编译不过，可以用纯 PyTorch 的 fallback 实现先跑通逻辑，但训练速度会慢三到五倍，只适合小规模验证。

还有一个坑：多卡训练时 selective scan 的并行实现有时会有 race condition，表现为 loss 偶尔跳变。遇到这种情况先退回单卡确认模型本身没问题，再检查多卡同步逻辑。

6. 和其他主干的横向对比实测

6.1 精度与显存的权衡表

我在一个内部肝脏分割数据集（约 800 例 CT，512×512）上做过对比，结果大致如下：

可以看到 VMamba 在精度和显存上都优于 Swin，但训练速度因为 selective scan 的串行特性，比纯卷积慢。单向版本精度掉得厉害，直接证明了四向扫描的价值。

6.2 什么场景该选 VMamba

不是所有医学分割都适合上 VMamba。我的判断标准是：如果目标结构需要长程上下文（大器官、多器官联合分割、需要全局解剖一致性），VMamba 优势明显；如果只是分割小而孤立的病灶（比如肺结节），卷积网络加足够的深度可能就够了，上 VMamba 反而增加训练成本。另外数据量太小时（几百张以下），VMamba 的参数量容易过拟合，建议先用预训练权重微调，或者退回轻量卷积网络。

6.3 预训练权重的使用

VMamba 官方提供了在 ImageNet 上预训练的分类权重，做分割时加载主干权重能显著加速收敛。但要注意：医学图像和自然图像的分布差异大，直接冻结主干训练解码器效果一般，最好是解冻全部做端到端微调，只是主干学习率调小。我试过只训解码器，Dice 比端到端低两个点左右。

7. 一些让 SS2D 更适配医学图像的改造思路

7.1 各向异性扫描

医学图像经常是各向异性的，比如 CT 的层厚方向分辨率远低于层内。如果做 3D 分割，把 SS2D 扩展到三个方向（加一个深度方向扫描）是自然的想法。但深度方向的序列长度通常很短（几十层），扫描收益有限，反而增加计算。我的做法是只在层内做四向扫描，层间用 3D 卷积处理，这样兼顾了效率和上下文。

7.2 多尺度 SS2D

单尺度扫描对大小差异大的目标不够友好。可以在解码器的不同阶段用不同的 d_state：浅层用小 d_state（比如 8）关注局部细节，深层用大 d_state（比如 32）建模全局。这样参数分配更合理，实测比统一 d_state 好一个点左右。

7.3 边界感知的扫描顺序

一个比较新的思路是让扫描顺序本身携带边界信息，比如在器官边界处改变扫描的起点，让边界像素在序列里更靠前，获得更强的状态更新。这个我还在实验阶段，初步结果有提升但不稳定，等验证充分了再单独写一篇。

整体做下来，VMamba 的 SS2D 四向扫描确实是医学图像分割里一个值得投入的方向，它把线性复杂度和长程建模这两个原本矛盾的需求捏到了一起。但它不是银弹，显存优化、损失设计、边界处理这些传统功夫一样都不能少。我个人的体会是，把它当成一个更强的主干来用，而不是指望它自动解决所有问题，心态就对了。后续如果大家在做 3D 或者多模态医学分割，SS2D 的扩展空间还很大，值得继续挖。