SwinUNETR架构自监督训练CVPR2022论文

这篇论文是 CVPR 2022 的论文，官方提出了 Swin UNETR 架构及其自监督预训练方法，这里给出了预训练权重。

📄 论文核心信息

项目	内容
标题	Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis
作者	Yucheng Tang, Dong Yang, Wenqi Li, Holger R. Roth, Ali Hatamizadeh 等 (NVIDIA + Vanderbilt)
发表	CVPR 2022
代码	https://monai.io/research/swin-unetr

---

1. 研究动机与核心贡献

为什么需要这个工作？

1. 领域鸿沟：自然图像（ImageNet）与医学图像（CT/MRI）存在巨大差异，直接迁移效果有限。
2. 3D数据标注困难：3D医学图像的专家标注极其昂贵且耗时，大规模有监督预训练不现实。
3. 2D Transformer的局限：将3D体积按切片处理会丢失跨平面的空间上下文信息。

核心贡献（三点）

1. 提出 Swin UNETR 架构：首个将 Swin Transformer 编码器与 CNN 解码器结合的 3D Transformer 分割模型。
2. 自监督预训练框架：设计了三个代理任务（掩码修复、旋转预测、对比学习）联合训练编码器。
3. SOTA 结果：在 BTCV（13器官）和 MSD（10任务）两个公开基准上取得排行榜第一。

---

2. Swin UNETR 架构细节

编码器（Swin Transformer）

配置项	论文设置
输入尺寸	96×96×96（预训练） / 128×128×128（微调可选）
Patch Size	2×2×2
Embedding Dimension	48
窗口大小 (Window Size)	7×7×7
Stage 数量	4 个 Stage
每个 Stage 的 Block 数	(2, 2, 2, 2) → 共 8 层
每层头数	(3, 6, 12, 24)

架构特点：

• 每个 Stage 包含 2 个 Swin Transformer Block（W-MSA + SW-MSA）
• Stage 间通过 Patch Merging 下采样（分辨率减半，通道翻倍）
• 输出 4 个不同尺度的特征图（1/2, 1/4, 1/8, 1/16）

解码器（CNN-based）

• 从编码器的 4 个 Stage 提取特征，通过跳跃连接传入 CNN 解码器
• 每个尺度先用 3×3×3 卷积处理，然后上采样并与前一尺度拼接
• 最终通过 1×1×1 卷积 + Softmax 输出分割概率图

---

3. 自监督预训练方法（核心）

三个代理任务

代理任务	目的	损失函数
掩码体素修复	学习纹理、结构和上下文对应关系	L1 Loss
3D 旋转预测	学习图像的几何结构内容	Cross-Entropy Loss
对比学习	最大化同一样本不同增广视图的互信息	InfoNCE Loss

训练流程

输入 CT 子体积 (96³)
    │
    ├─→ 随机裁剪 + 旋转 (0°, 90°, 180°, 270°)
    │        │
    │        └─→ 生成两个增广视图 (正样本对)
    │
    ├─→ 随机 Cutout (掩码比例 30%)
    │
    ├─→ Swin UNETR 编码器
    │        │
    │        ├─→ 修复头 → L1 重建损失
    │        ├─→ 旋转头 → 交叉熵分类损失
    │        └─→ 对比头 → InfoNCE 对比损失
    │
    └─→ 总损失 = Linpaint + Lcontrast + Lrot (λ₁=λ₂=λ₃=1)

预训练超参数

参数	值
学习率	4e-4 (AdamW + 余弦退火)
Batch Size	4 per GPU
输入尺寸	96×96×96
迭代次数	450K steps
训练数据	5,050 张 CT 图像（训练集 4,761，验证集 249）
GPU 天数	约 6 GPU 天

---

4. 预训练数据详情

数据来源

数据集	部位	样本数
LUNA16	胸部	888
TCIA COVID-19	胸部	761
LIDC	胸部	475
HNSCC	头颈部	1,287
TCIA Colon	腹部/盆腔	1,599
合计	头、颈、胸、腹、盆腔	5,050

预处理

• CT 值截断到 [-1000, 1000] → 归一化到 [0,1]
• 只裁剪包含前景（非空气）的 96×96×96 子体积
• 不使用任何标注信息

---

5. 实验结果

BTCV（13器官分割）

方法	平均 Dice
nnUNet	0.888
UNETR	0.891
Swin UNETR (预训练)	0.918

关键发现：

• 小器官提升明显：胰腺 +1.6%，肾上腺 +3.8%，门静脉 +3.6%
• 用预训练权重仅需 60% 标注数据 即可达到从头训练全量数据的性能

MSD（10任务）

方法	平均 Dice	排名
nnUNet	77.89	第3
DiNTS	77.93	第2
Swin UNETR (预训练)	78.68	第1

---

6. 关键结论

1. 预训练有效：在所有 CT 任务（肝脏、肺、胰腺、肝血管、脾脏、结肠）上，预训练模型均优于从头训练。
2. 小样本优势明显：仅用 10% 标注数据，预训练模型比从头训练高约 10 个 Dice 点。
3. 数据规模影响：预训练数据越多，微调性能越好（100→3000→5000 扫描持续提升）。
4. CT → MRI 迁移有限：因模态差异，预训练权重对 MRI 任务帮助不大（实验中对 MRI 任务从头训练）。

---

7. 对使用预训练权重的建议

预训练权重的特性

特性	说明
包含内容	仅编码器（Encoder）权重，不包含解码器
输入尺寸	预训练时使用 96×96×96，固定了位置编码
最佳使用方式	在 96³ 尺寸上微调，或使用 load_from 加载（自动忽略位置编码不匹配）

论文中的微调设置（参考）

数据集	输入尺寸	学习率	数据增强
BTCV	96×96×96	1e-4	随机翻转、旋转、强度偏移
MSD	96×96×96	1e-4~1e-3	任务相关

8. MONAI社区扩展

这是扩展版预训练权重的直接来源。教程中明确说明预训练使用了多个公开数据集，包含超过50,000个3D CT和MR体积。

资源链接: Project-MONAI/tutorials - swinunetr_pretrained

官方在 BTCV (Beyond the Cranial Vault) 数据集上进行了微调实验，结果清楚地展示了使用预训练权重的巨大优势。效果如下表。

即使在仅有 3 个训练样本的极端情况下，使用预训练权重的模型 Dice 分数也达到了 60.47，显著高于随机初始化的 58.29。这充分证明了该权重在小样本学习任务中的巨大价值。

Training Volumes	% of Dataset	Validation Volumes	Random Init Dice score	Pre-trained Dice Score
3	10%	6	58.29	60.47
5	20%	6	67.46	68.25
7	30%	6	71.22	72.43
12	50%	6	79.39	79.94
24	100%	6	82.63	82.64