Swin-UNETR医学影像分割编码器

可视化演示

展示Swin-UNETR三大核心机制：窗口注意力(W-MSA)在局部窗口内计算自注意力、移位窗口(SW-MSA)实现跨窗口信息交互、编码器-解码器架构通过跳跃连接融合多尺度特征进行3D医学影像分割，动态展示Dice分数和注意力覆盖率。

分割任务概述

你站在教堂中央，抬头仰视巨大的穹顶壁画。修复委托人递给你一张表格，上面列出了图中所包含的所有构件：圣徒衣袍、天使翅膀、云朵阴影、石柱纹理。你的第一项工作，就是把扫描仪生成的那堆立体小方块，逐块核对它们属于表格中的哪一类。在医学影像AI里，这份「逐块核对」的工作就是语义分割——只不过对象换成了病人的CT或MRI扫描，而「构件表格」换成了器官清单，比如肝脏、脾脏、肿瘤、血管。

CT和MRI是医学影像的两大主力。CT扫描像给身体拍了一张「密度地图」，骨骼和金属植入物在图像里闪闪发亮；MRI则更像「软组织对比图」，对肌肉、脂肪、脑组织的变化非常敏感。它们输出的数据，本质上都是一个三维数组——沿着身体的长、宽、厚三个方向均匀采样，得到一串串离散的小立方体。每个小立方体就是「体素」（三维空间的像素），里面存着一个数字，代表那一点组织的信号强度。

在数学上，分割任务可以被定义为一个逐体素分类问题。如果影像有 H×W×D 个体素，模型要做的事就是为每一个体素输出一个类别编号。这种「像素级精细标注」的能力，正是肿瘤定位、手术规划和放疗靶区勾画等临床场景的基石。

回到壁画场景：编号就是你在修复图纸上使用的颜色代码——编号0留白，编号1画上金色衣袍，编号2画上白色翅膀。分割模型的全部努力，就是让这个编号足够精准。

编码器架构

拿到了壁画的三维扫描数据之后，修复师团队开始分工。团队负责人没有让所有人同时盯着整幅壁画——那会让注意力过于分散，而是派出四组观察员，每组负责一个尺度层级，每一组内部又用「窗口」把壁画切成小块轮流细看。这套分工体系，就是Swin Transformer编码器。

编码器的输入端，首先把原始三维影像切成不重叠的小立方体块（patch）。Swin-UNETR使用的patch尺寸是2×2×2，也就是说，每个2×2×2的小块被压成一个「令牌」（token）——像是把一小团马赛克捏成一个代表符。接着，一个线性层把这个令牌从低维映射到高维特征空间，旅程正式启程。

随后，编码器经过四个stage，每个stage由若干Swin Block组成。每个stage结束后，特征图的分辨率减半、通道数翻倍：从原始96³体素，一路走到1/4、1/8、1/16、1/32分辨率，通道数也从48逐级涨到768。这正是Swin-UNETR能够同时兼顾「大范围上下文」和「局部细节」的骨架所在。

观察员们拿到了各自的「小块笔记」后，并不会立刻交差——他们还要经历接下来几道工序，把笔记逐步整合成有层次的结构理解。而这套工序的核心，就是注意力机制。

注意力机制

观察员领到的聚光灯，是Swin-UNETR最核心的工具。它一次只照亮壁画的一小块区域（比如2×2×2个令牌组成的窗口），在这个窗口内部，让每个令牌与窗口内其他令牌互相「打招呼」——互相确认彼此的笔触、颜色、纹理，判断自己与它们的关系。这种窗口内的自注意力，就是「W-MSA」。它像修复师在灯下眯着眼，把一小片壁画的每一笔都与其他每一笔做比对。

但只盯着窗口内部有个致命问题：窗口与窗口之间永远不会互相看到，壁画上横跨窗口边界的长线条（比如天使手中横贯壁画的长杖）就会被拦腰截断。于是Swin Transformer引入了「移动窗口」（SW-MSA）：在下一层，所有窗口整体向右下偏移半个窗口的距离，让原先的边界变成了新窗口的内部。这样一来，信息就能跨窗口流动，聚光灯也照亮了衔接处。

为了让模型知道「窗口内谁和谁离得近」这个空间信息，Swin又在注意力计算里加入了相对位置偏置B。偏置好比是修复师的经验：笔触越靠近，越可能属于同一片衣袍；相隔越远，越可能是不同的装饰纹样。

一盏是「普通窗口聚光灯」，一盏是「移动后的聚光灯」，两盏灯在相邻的神经层里交替使用。壁画上再长的线条，也逃不过它们的接力照亮。

层级特征提取

观察员的观察笔记越记越厚。负责人开始要求他们「退后几步看全貌」：把2×2×2个相邻令牌的笔记合并成一份「概括笔记」——分辨率缩一半，但每个概括笔记携带的信息维度翻一倍。这个过程叫Patch Merging，就像是把四张详细的局部素描叠影成一张更宏观的构图稿。

经过四轮这样的退后，编码器手里握有了四种尺度的笔记：最细的1/4尺度笔记保留了大量边缘和纹理；中间的1/8、1/16尺度开始出现「这是一片翅膀」这样的语义块；最粗的1/32尺度则几乎是一张完整的构图缩略图，每一个令牌都能「看到」几乎整幅壁画。这四份笔记合在一起，就是计算机视觉里常说的「特征金字塔」。

每一层的语义也发生着质变：浅层特征回答「这个笔触是红色还是蓝色」，深层特征回答「这块区域属于天使还是云朵」。修复工作最怕「只缘身在此山中」——只盯着局部会迷失方向，只看全局又看不清裂纹。而Swin的多尺度笔记，正好让团队既有全局方位感，又不放过任何细节。

当壁画的结构笔记和细节笔记都齐备之后，下一步就是把分散在四位观察员手里的笔记，交回给总修复师做最终拼装。

解码器结合

总修复师面前摆开了两份材料：一份是观察员们交来的多尺度笔记，一份是当初的三维扫描原始数据。他需要的是一张完整的修复图纸——和壁画同样大小的标签图，而不是一堆分散的笔记。于是，U-Net解码器登场了。

解码器的工作方式是「由粗到细逐级还原」：最粗的1/32尺度笔记先通过反卷积上采样放大一倍，再与编码器同层的1/16尺度笔记拼接在一起。这里的拼接就是「跳跃连接」（Skip Connection）——它像一份「原始草图」，用编码器侧的高分辨率特征为解码器校准位置、补回细节纹理。每一级都重复这个「放大→拼接→融合」的过程，直到特征图恢复到原始分辨率。

最后，一个1×1×1卷积把特征映射成每个体素在K个类别上的得分，SoftMax之后就是最终的分割概率图。为了方便梯度传导，Swin-UNETR还会在多个解码层级同时计算损失，这种「深监督」策略就像师傅在每一道工序都抽查一次质量，而不是只验收最后成品。

总修复师捧着拼接好的图纸走完四个还原阶段，一张跟原始壁画同样尺寸、每一个马赛克都有明确归属的修复图纸终于成型。剩下的事情，就是让整支团队接受考核，不断打磨手艺。

训练优化

修复师团队要接受严格考核。考官手里也有一份专家绘制的标准答案图纸，团队画出的图纸与标准答案重叠得越多，得分越高。这个「重叠率」在医学影像分割里就是Dice系数——你预测的器官区域与医生标注的真实区域之间的重合程度。完全重合是1分，毫无交集是0分。

Swin-UNETR训练时最常用的损失函数是Dice损失与交叉熵损失的组合。Dice损失从「区域重叠」角度打分，交叉熵从「每个体素是否归类正确」角度打分，二者互补——前者关心完整性，后者关心精确性。配合AdamW优化器、余弦退火学习率，模型在BraTS数据集上通常训练数万步即收敛。

数据增强方面，团队采用了「随机裁剪96³小块、随机翻转、随机旋转、灰度缩放」等策略。这就像让修复师们反复练习不同光线、不同角度下的壁画局部，练得多了，自然不怕真实世界的千变万化。最终评估时，Dice和HD95（表面距离）是公认的两大硬指标。

考场的模拟训练日复一日。当团队在验证集上的Dice分数稳定不再上升时，负责人一声令下：上真实战场。

应用效果

修复师团队终于走出集训营，迎来真实教堂的验收。医学影像界的「验收现场」是BraTS（脑肿瘤分割挑战赛）——一个汇聚全球顶尖算法的年度盛事，提供带专家标注的多模态MRI数据。Swin-UNETR一登场就交出了亮眼成绩：在BraTS 2021验证集上，整体肿瘤Dice达到约0.847，超越原来的UNETR约2个百分点，跻身榜单前列。

这项成绩很大程度上要归功于「多模态输入」的设计。BraTS为每个病人提供了四种MRI模态——T1、T1ce（增强）、T2和FLAIR，Swin-UNETR将四种模态按通道方向拼接成4通道输入，就像修复师既看素描底稿又看彩色照片，能捕捉到单一模态下难以察觉的病灶差异。这种多模态融合能力，也延续到CT多期增强扫描等场景中。

从比赛到临床，Swin-UNETR的价值正在快速落地。放射科医生手动勾画一个肝脏肿瘤通常需要十几分钟，而模型推理只需数秒；在放疗靶区勾画、脑肿瘤随访、器官移植体积评估等任务中，它都能充当医生的「第二双眼睛」。当然，临床部署还要经过严谨的验证与审批，但技术的齿轮已经开始转动。

壁画修复师团队捧回了奖杯，但这远不是终点。医院影像科里还有无数幅「壁画」等待标注——每一次Swin-UNETR的部署，都是对「分块细看、逐级拼装」这套方法论的一次实战检验。

AI 应用场景

Python 演示

import torchimport torch.nn as nn  class SwinBlock3D(nn.Module):    """简化版 3D Swin Transformer Block (W-MSA + SW-MSA 核心机制)"""     def __init__(self, dim, window_size=2, num_heads=2):        super().__init__()        self.window_size = window_size        self.num_heads = num_heads        self.norm1 = nn.LayerNorm(dim)        self.norm2 = nn.LayerNorm(dim)        self.qkv = nn.Linear(dim, dim * 3)        self.mlp = nn.Sequential(            nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim)        )     def _window_partition(self, x, w):        B, D, H, W, C = x.shape        x = x.view(B, D // w, w, H // w, w, W // w, w, C)        x = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous()        return x.view(-1, w, w, w, C)     def _window_reverse(self, windows, w, D, H, W):        B = windows.shape[0] // (D // w * H // w * W // w)        C = windows.shape[-1]        x = windows.view(B, D // w, H // w, W // w, w, w, w, C)        x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous()        return x.view(B, D, H, W, C)     def _window_attn(self, x):        B, D, H, W, C = x.shape        windows = self._window_partition(x, self.window_size)        Bw, wD, wH, wW, Cw = windows.shape        tokens = windows.reshape(Bw, wD * wH * wW, Cw)        qkv = self.qkv(tokens)        qkv = qkv.reshape(Bw, wD * wH * wW, 3, self.num_heads, C // self.num_heads)        qkv = qkv.permute(2, 0, 3, 1, 4)        q, k, v = qkv[0], qkv[1], qkv[2]        attn = (q @ k.transpose(-2, -1)) * (C // self.num_heads) ** -0.5        attn = attn.softmax(dim=-1)        out = (attn @ v).transpose(1, 2).reshape(Bw, wD, wH, wW, Cw)        return self._window_reverse(out, self.window_size, D, H, W)     def forward(self, x):        B, D, H, W, C = x.shape        shortcut = x        x = self.norm1(x)        x = shortcut + self._window_attn(x)  # W-MSA 窗口内注意力        shortcut = x        x = torch.roll(x, shifts=(1, 1, 1), dims=(1, 2, 3))  # 循环移位        x = self.norm2(self._window_attn(x))        x = torch.roll(x, shifts=(-1, -1, -1), dims=(1, 2, 3))  # 还原位置        return shortcut + x  # SW-MSA 移动窗口注意力  if __name__ == "__main__":    block = SwinBlock3D(dim=16, window_size=2, num_heads=2)    x = torch.randn(1, 4, 4, 4, 16)    y = block(x)    assert tuple(y.shape) == tuple(x.shape)    print("输入形状:", tuple(x.shape))    print("输出形状:", tuple(y.shape))    print("Swin Block 前向传播成功: 形状不变, 注意力计算正常")

1. self._window_partition(x, self.window_size)把特征图切成多个不重叠的窗口，就像修复师把壁画分成2×2×2的小区域。
2. torch.roll(x, shifts=(1, 1, 1), dims=(1, 2, 3))对三个空间维度做循环移位，让窗口边界移动半个窗口，实现SW-MSA。
3. qkv = qkv.reshape(Bw, wD * wH * wW, 3, self.num_heads, C // self.num_heads)把令牌映射成多头的Q、K、V，聚光灯开始逐头扫描。
4. attn = (q @ k.transpose(-2, -1)) * (C // self.num_heads) ** -0.5点积打分并除以√d，得到窗口内令牌两两的相关性。
5. ctx.fillRect(ox + pos.x * cell, oy + pos.y * cell, WIN * cell, WIN * cell)Canvas里高亮当前移动窗口（呼应代码中的窗口划分）。

常见误区

面试题