介绍 Swin UNETR 模型。

3 Swin UNETR

3.1 编码器（Encoder）

我们在图 1 中展示了 Swin UNETR 的架构。Swin UNETR 模型的输入 \(\mathcal{X} \in \mathbb{R}^{H \times W \times D \times S}\) 是一个张量，其块分辨率为 \((H', W', D')\)，维度为 \(H' \times W' \times D' \times S\)。我们首先利用一个 Patch Partition（块划分）层，创建一个 3D 令牌序列，其维度为 \(\left[\frac{H}{H'}\right] \times \left[\frac{W}{W'}\right] \times \left[\frac{D}{D'}\right]\)，并将它们投影到维度为 \(C\) 的嵌入空间中。自注意力在划分阶段创建的非重叠窗口内计算，以实现高效的令牌交互建模。

图 2 展示了后续层的移位窗口机制。具体来说，在 Transformer 编码器的第 \(l\) 层，我们使用大小为 \(M \times M \times M\) 的窗口，将一个 3D 令牌均匀划分为 \(\left[\frac{H'}{M}\right] \times \left[\frac{W'}{M}\right] \times \left[\frac{D'}{M}\right]\) 个区域。随后，在第 \(l+1\) 层，划分的窗口区域沿各维度移动 \(\left(\left[\frac{M}{2}\right], \left[\frac{M}{2}\right], \left[\frac{M}{2}\right]\right)\) 个体素。在后续层中，这种移位窗口机制交替进行，从而使得相邻窗口之间能够交互信息。

3.2 损失函数（Loss Function）

我们使用 软 Dice 损失函数（soft Dice loss） [30]，它以体素方式计算，公式如下：

\[ \mathcal{L}(G,Y) = 1 - \frac{2}{J}\sum_{j = 1}^{J}\frac{\sum_{i = 1}^{I}G_{i,j}Y_{i,j}}{\sum_{i = 1}^{I}G_{i,j}^{2} + \sum_{i = 1}^{I}Y_{i,j}^{2}}. \quad (3) \]

其中：

• \(I\) 表示体素总数；
• \(J\) 表示类别数；
• \(G_{i,j}\)：第 \(i\) 个体素在第 \(j\) 类上的 真实标签（Ground Truth），值为 0 或 1（one-hot 编码）
• \(Y_{i,j}\)：第 \(i\) 个体素在第 \(j\) 类上的 预测概率，值为 0~1 之间的连续值（经过 softmax 输出）

计算流程（以单个类别 \(j\) 为例）

1. 分子：\(\sum_{i} G_{i,j} Y_{i,j}\) —— 真实标签与预测概率的逐元素乘积之和，即预测正确的部分（交集）。
2. 分母：\(\sum_{i} G_{i,j}^{2} + \sum_{i} Y_{i,j}^{2}\) —— 真实标签的平方和与预测概率的平方和之和（相当于两集合的模长平方之和）。
3. 比值：\(2 \times \frac{\text{分子}}{\text{分母}}\) —— 即 Dice 系数，取值范围 0~1，值越大表示分割越准确。
4. 损失：\(1 - \text{Dice系数}\) —— 转化为损失函数，值越小越好。

为什么叫“软”Dice？

• 硬 Dice：使用 0/1 二值预测（即取最大类别），适合离散输出。
• 软 Dice：使用 连续概率值（如 softmax 输出），允许网络在训练过程中进行梯度回传和反向传播，更适合端到端训练。

核心优势

Dice 损失天然适合处理类别极度不平衡的分割任务（如肿瘤只占图像很小一部分）。因为它直接优化 重叠度，而不像交叉熵那样对每个像素独立考虑，从而能有效缓解背景主导梯度的问题。

3.3 实现细节（Implementation Details）

Swin UNETR 使用 PyTorch³ 和 MONAI⁴ 实现，并在配备 8 块 NVIDIA V100 GPU 的 DGX-1 集群上训练。表 1 详细列出了 Swin UNETR 架构的配置、参数量和 FLOPs。

学习率设置为 0.0008。我们根据非零体素将所有输入图像归一化为零均值和单位标准差。训练时，从 3D 图像体积中随机裁剪 \(128 \times 128 \times 128\) 的块。我们以 0.5 的概率对所有三个轴应用随机轴镜像翻转。此外，我们还应用随机强度偏移（范围为 \([-0.1, 0.1]\)）和随机强度缩放（范围为 \([0.9, 1.1]\)）。

表 1. Swin UNETR 配置

嵌入维度	特征尺寸	块数量（各阶段）	窗口大小	注意力头数（各阶段）	参数量	FLOPs
768	48	[2, 2, 2, 2]	[7, 7, 7]	[3, 6, 12, 24]	61.98M	394.84G

嵌入维度-与标准swin不同点

• 特征尺寸（feature size）: Patch Embedding 后映射到 48 维的向量空间，就是 SWIN 中的 C 。当数据从一个阶段传到下一个阶段（Patch Merging），空间分辨率减半，但“员工人数”（通道数）会翻倍：48 → 96 → 192 → 384。
• 嵌入维度（Embedding Dimension）：Transformer 内部每个 token 的特征向量长度，具体做法如下，这就是和标准的 swin transformer 不一样的地方。

具体计算过程：

1. 假设输入一个 48 维的向量 \(X\)（形状 [48]）（或后续阶段的 96、192、384）。
2. 计算 Q、K、V 时，会使用三个矩阵 \(W_Q, W_K, W_V\)（形状为 [48, 768]）进行乘法：
   • \(Q = X \cdot W_Q\)，得到形状 [768]。
   • 同理 K、V 也变为 768 维。
3. 在 768 维空间里计算点积注意力（\(Q \cdot K^T\)）。
4. 注意力结果会通过另一个线性层 \(W_O\)（形状 [768, 48]），映射回 48 维，作为该块的输出。

3.4 数据集与模型集成（Dataset and Model Ensembling）

BraTS 挑战赛旨在通过提供带有医生标注的体素级真实标签的 3D MRI 数据集，来评估脑肿瘤语义分割的最先进方法 [6, 29, 5, 3, 4]。BraTS 2021 挑战赛的训练数据集包含 1251 例受试者，每例包含四种 3D MRI 模态：a) 原生 T1，b) 增强后 T1 加权（T1Gd），c) T2 加权（T2），以及 d) T2 液体衰减反转恢复（T2-FLAIR）。这些数据经过刚体对齐，并重采样至 \(1 \times 1 \times 1\) mm 的各向同性分辨率，且已完成去颅骨处理。输入图像尺寸为 \(240 \times 240 \times 155\)。数据收集自多个机构，使用多种 MRI 扫描仪。

标注包括三个肿瘤子区域：增强肿瘤（ET）、瘤周水肿（ED） 以及坏死和非增强肿瘤核心（NCR/NET）。这些标注被组合为三个嵌套的子区域：全肿瘤（WT）、肿瘤核心（TC） 和 增强肿瘤（ET）。图 3 展示了所有语义类别的典型分割输出。

在本次挑战赛中，还提供了两个不包含真实标签的额外数据集，分别用于验证阶段和测试阶段。这些数据集要求参与者将分割掩膜上传至主办方服务器进行评估。验证数据集包含 219 例病例，用于中间模型评估。测试数据集的相关信息未向参与者提供。

我们的模型在 BraTS 2021 数据集上训练，训练集和验证集分别包含 1251 例和 219 例病例。验证集对应的分割标签未公开，性能基准通过向 BraTS 2021 官方服务器提交获得。我们采用 五折交叉验证，划分比例为 80:20。我们未使用任何额外数据。

最终结果通过对 10 个 Swin UNETR 模型进行集成获得，以提高性能并对所有预测达成更好的一致性。这些集成模型来自两次独立的五折交叉验证训练。

如何集成

他们完整运行了两轮五折交叉验证（可能使用不同随机种子或数据划分），每轮产生 5 个模型，两轮共得到 10 个模型。

在推理时，将 同一个输入图像 分别输入这 10 个模型，得到 10 个分割结果（每个像素属于某类别的概率）。最终结果通过 投票或平均概率 决定：

• 概率平均：对 10 个模型的 softmax 输出求均值，取最大概率的类别作为最终预测。
• 这种方法能有效 平滑单个模型的偶然误差，提升整体鲁棒性和 Dice 分数。

为什么有效？

• 不同模型在训练时看到了不同的数据子集，学到了略有差异的决策边界。
• 集成可以减少过拟合风险，尤其是在数据有限时。
• 在 BraTS 这种高竞争性挑战中，集成的 0.5%~1% Dice 提升往往是胜负关键。

论文中没有介绍解码器，这里讨论一下。

解码器

虽然论文对解码器的描述比较简略，但结合其官方实现（MONAI）和相关资料，可以清晰地还原出它的结构。SwinUNETR的解码器是一个基于CNN的、与编码器结构对称的上采样网络，负责将抽象特征恢复为精细的分割图。

🧬 解码器的核心输入：跳跃连接

解码器并非只接收编码器最深层的输出，它的输入由两部分组成：

1. 来自编码器对应层的特征：编码器（Swin Transformer）在4个不同分辨率阶段输出的特征图，会通过跳跃连接（Skip Connections） 直接传递到解码器的对应阶段。

   看架构图实际是5个，除了4个stage的输出，还包括了 Patch Embedding层的输出，其分辨率为 \(H/2 \times W/2 \times D/2\)，之后分辨率逐级递减，并且通道数目翻倍

2. 来自上一解码器层上采样后的特征：解码器自身深层（分辨率更低）的输出，会通过上采样逐步恢复到更高分辨率。

🏗️ 解码器的核心组件与工作流程

解码器的主体由4个上采样阶段组成，与编码器的4个阶段一一对应。每个阶段的核心操作是：

1. 特征重整形与残差处理 (Reshape & Residual Block)：来自编码器的跳跃连接特征，会先被重整形，然后通过一个残差块（Residual Block） 进行处理。这个残差块通常包含两个 3×3×3 的3D卷积层，每层后都跟随实例归一化（Instance Normalization） 和激活函数。
2. 上采样 (Upsampling)：经过残差块处理后的特征图，会通过转置卷积（Transposed Convolution） 或插值+卷积的方式进行上采样，将空间分辨率提高一倍。
3. 特征拼接与融合 (Concatenation & Refinement)：将上采样得到的特征图，与来自编码器同层（通过跳跃连接） 的特征图在通道维度上拼接（Concatenate） 起来。融合后的特征会再通过一个残差块进行 refinement。

🎯 最后的输出层

经过4个上采样阶段后，解码器最终会输出一个与原始输入图像空间分辨率相同的特征图。最后，通过一个 1×1×1 的卷积层（UnetOutBlock），将通道数映射到目标类别数量（如3），并用Sigmoid或Softmax函数输出最终的分割概率图。

💎 总结

SwinUNETR的解码器是一个标准的、基于CNN的U-Net解码器。它的设计简洁而高效，通过跳跃连接融合编码器的多尺度特征，并利用残差块和上采样逐步恢复空间细节，最终生成精准的分割结果。你可以参考MONAI官方文档中关于 UnetrUpBlock 的源码来了解更具体的实现细节。