Swin-Unet

[ECCVW2022] 本项目包含论文《Swin-Unet：用于医学图像分割的类Unet纯Transformer》（https://arxiv.org/abs/2105.05537）的代码。我们的论文已被ECCV 2022 医学计算机视觉研讨会（https://mcv-workshop.github.io/）接收。我们已更新了可复现性说明，希望这能帮助您复现实验结果。

1. 下载预训练的Swin Transformer模型（Swin-T）

• [在此链接获取预训练模型] (https://drive.google.com/drive/folders/1UC3XOoezeum0uck4KBVGa8osahs6rKUY?usp=sharing)：将预训练的Swin-T模型放入“pretrained_ckpt/”文件夹中。

2. 准备数据

• 我们使用的数据集由TransUnet的作者提供。[在此链接获取处理好的数据]（Synapse/BTCV：https://drive.google.com/drive/folders/1ACJEoTp-uqfFJ73qS3eUObQh52nGuzCd；ACDC：https://drive.google.com/drive/folders/1KQcrci7aKsYZi1hQoZ3T3QUtcy7b--n4）。

3. 环境配置

• 请准备一个Python版本为3.7的环境，然后使用命令“pip install -r requirements.txt”安装依赖项。

4. 训练与测试

• 在Synapse数据集上运行训练脚本。我们使用的批量大小为24。如果您显存不足，可以将批量大小降低至12或6以节省显存。

• 训练

sh train.sh 
# 或 
python train.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path your DATA_DIR --max_epochs 150 --output_dir your OUT_DIR  --img_size 224 --base_lr 0.05 --batch_size 24

• 测试

sh test.sh 
# 或 
python test.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --is_saveni --volume_path your DATA_DIR --output_dir your OUT_DIR --max_epoch 150 --base_lr 0.05 --img_size 224 --batch_size 24

可复现性

• 代码

我们训练好的模型存储在华为云上。实习生无权将任何文件从内部系统导出，因此我无法分享我们的模型权重。关于如何复现论文中展示的分割结果，我们发现不同类型的GPU会产生不同的结果。在我们的代码中，我们仔细设置了随机种子，因此在同一类型GPU上多次训练时，结果应当一致。如果训练得到的分割结果与论文中的不一致，建议调整学习率。此外，我们在这项工作中使用的GPU类型是Tesla V100。最后，对于纯Transformer模型而言，预训练非常重要。在我们的实验中，编码器和解码器都使用了预训练权重进行初始化，而不仅仅是编码器使用预训练权重。

参考文献

• TransUnet
• SwinTransformer

引用

@InProceedings{swinunet,
author = {Hu Cao and Yueyue Wang and Joy Chen and Dongsheng Jiang and Xiaopeng Zhang and Qi Tian and Manning Wang},
title = {Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation},
booktitle = {欧洲计算机视觉会议研讨会论文集（ECCVW）},
year = {2022}
}

@misc{cao2021swinunet,
      title={Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation}, 
      author={Hu Cao and Yueyue Wang and Joy Chen and Dongsheng Jiang and Xiaopeng Zhang and Qi Tian and Manning Wang},
      year={2021},
      eprint={2105.05537},
      archivePrefix={arXiv},
      primaryClass={eess.IV}
}

Swin-Unet 快速上手指南

Swin-Unet 是一种基于纯 Transformer 架构的 U-Net 变体，专为医学图像分割任务设计。本指南将帮助您快速搭建环境并运行模型。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python 版本: 3.7
• 硬件要求: NVIDIA GPU (论文实验基于 Tesla V100，显存不足时可调整 batch size)
• 前置依赖: 需安装 requirements.txt 中列出的依赖包。

2. 安装步骤

2.1 创建虚拟环境并安装依赖

建议使用 Conda 创建独立的 Python 3.7 环境，然后安装项目依赖。

conda create -n swinunet python=3.7
conda activate swinunet
pip install -r requirements.txt

2.2 下载预训练模型

下载 Swin-Tiny 预训练权重，并将其放入项目根目录下的 pretrained_ckpt/ 文件夹中。

• 下载地址: Google Drive
• 注：国内用户若无法访问 Google Drive，请自行寻找对应的 Swin-Tiny ImageNet 预训练权重文件。

2.3 准备数据集

本项目使用的数据集（Synapse/BTCV 和 ACDC）由 TransUnet 作者提供。请下载处理好的数据并解压。

• Synapse/BTCV: 下载链接
• ACDC: 下载链接

3. 基本使用

以下以 Synapse 数据集为例，展示训练和测试的最简命令。请根据实际情况替换 your DATA_DIR（数据路径）和 your OUT_DIR（输出路径）。

3.1 训练模型

默认批量大小（batch size）为 24。若显存不足，可将 --batch_size 调整为 12 或 6。

python train.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path your DATA_DIR --max_epochs 150 --output_dir your OUT_DIR  --img_size 224 --base_lr 0.05 --batch_size 24

或者使用脚本运行：

sh train.sh 

3.2 测试模型

使用训练好的模型进行推理和评估。

python test.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --is_saveni --volume_path your DATA_DIR --output_dir your OUT_DIR --max_epoch 150 --base_lr 0.05 --img_size 224 --batch_size 24

或者使用脚本运行：

sh test.sh 

复现提示：

1. 纯 Transformer 模型对预训练权重非常敏感，本代码中编码器和解码器均使用了预训练权重初始化。
2. 不同 GPU 型号可能导致结果细微差异，若结果与论文不符，建议微调学习率（--base_lr）。
3. 代码已设置随机种子，在同类型 GPU 上多次训练结果应保持一致。