[AAAI2022] UCTransNet

这个仓库是论文《UCTransNet：从通道角度重新思考U-Net中的跳跃连接与Transformer》的官方实现，该论文已被AAAI2022接收。

框架图

我们提出了一个通道Transformer模块（CTrans），并用它替换原始U-Net中的跳跃连接，因此我们将其命名为"U-CTrans-Net"。

在线演示视频可用于简要介绍。

🔥🔥🔥 UCTransNet的改进版本请参考UDTransNet（通过可学习的跳跃连接缩小U-Net的语义差距：以医学图像分割为例），它实现了更高的性能和更低的计算成本。🔥🔥🔥

环境要求

使用以下命令从requirements.txt安装：

pip install -r requirements.txt

使用方法

注意：如果你在使用代码时遇到问题，issues可能会有帮助。

1. 数据准备

1.1. GlaS和MoNuSeg数据集

原始数据可以从以下链接下载：

MoNuSeg数据集 - 链接（原始）
GLAS数据集 - 链接（原始）

然后按以下格式准备数据集，以便轻松使用代码：

├── datasets
    ├── GlaS
    │   ├── Test_Folder
    │   │   ├── img
    │   │   └── labelcol
    │   ├── Train_Folder
    │   │   ├── img
    │   │   └── labelcol
    │   └── Val_Folder
    │       ├── img
    │       └── labelcol
    └── MoNuSeg
        ├── Test_Folder
        │   ├── img
        │   └── labelcol
        ├── Train_Folder
        │   ├── img
        │   └── labelcol
        └── Val_Folder
            ├── img
            └── labelcol

1.2. Synapse数据集

我们使用的Synapse数据集由TransUNet的作者提供。详情请参阅https://github.com/Beckschen/TransUNet/blob/main/datasets/README.md。

（可选）🔥🔥 使用自定义数据集

如果你想在自定义数据集上实现UCTransNet，最简单的方法是按照上述GlaS的方式组织文件结构。
确保图像为.jpg格式，掩码ID应与图像ID匹配，但扩展名为.png。
文件结构或命名约定的任何不一致都可能导致I/O错误。

2. 训练

如论文中提到的，我们引入了两种策略来优化UCTransNet。

第一步是更改Config.py中的设置，所有配置包括学习率、批量大小等都在其中。

2.1 联合训练

我们使用单一损失函数同时优化U-Net中的卷积参数和CTrans参数。运行：

python train_model.py

2.2 预训练

我们的方法只是替换了U-Net中的跳跃连接，因此U-Net中的参数可以用作预训练权重的一部分。

通过首先使用/nets/UNet.py训练一个经典的U-Net，然后使用预训练权重来训练UCTransNet，CTrans模块可以获得更好的初始特征。这种策略可以提高收敛速度，在某些情况下可能会改善最终的分割性能。

3. 测试

3.1. 获取预训练模型

在此，我们提供了在GlaS和MoNuSeg数据集上预训练的权重，如果你不想自己训练模型，可以通过以下链接下载：

3.2. 测试模型并可视化分割结果

首先，在Config.py中更改会话名称，与训练阶段相同。然后运行：

python test_model.py

你可以得到Dice和IoU分数以及可视化结果。

🔥🔥 Synapse数据集中所有类别的测试结果可以通过此链接下载。 🔥🔥

4. 可复现性

在我们的代码中，我们仔细设置了随机种子并将cudnn设置为"确定性"模式以消除随机性。然而，仍然存在一些可能导致不同训练结果的因素，例如cuda版本、GPU类型、GPU数量等。我们实验中使用的GPU是NVIDIA A40（48G），cuda版本为11.2。

特别是在多GPU情况下，上采样操作在随机性方面存在很大问题。更多详情请参见https://pytorch.org/docs/stable/notes/randomness.html。

在训练时，我们建议训练模型两次以验证是否消除了随机性。由于我们使用了早停策略，最终性能可能会因随机性而发生显著变化。

参考文献

UNet++: https://github.com/qubvel/segmentation_models.pytorch
Attention U-Net: https://github.com/bigmb/Unet-Segmentation-Pytorch-Nest-of-Unets
MultiResUNet: https://github.com/makifozkanoglu/MultiResUNet-PyTorch
TransUNet: https://github.com/Beckschen/TransUNet
Swin-Unet: https://github.com/HuCaoFighting/Swin-Unet
MedT: https://github.com/jeya-maria-jose/Medical-Transformer

引用

如果此代码对你的研究有帮助，请引用：

@article{UCTransNet,
	 title={UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-Wise Perspective with Transformer}, 
	 volume={36}, 
	 url={https://ojs.aaai.org/index.php/AAAI/article/view/20144}, 
  	 DOI={10.1609/aaai.v36i3.20144},
	 number={3}, 
	 journal={Proceedings of the AAAI Conference on Artificial Intelligence}, 
	 author={Wang, Haonan and Cao, Peng and Wang, Jiaqi and Zaiane, Osmar R.}, 
	 year={2022}, 
	 month={Jun.}, 
	 pages={2441-2449}}

联系方式

Haonan Wang (haonan1wang@gmail.com)