在 AutoDL 平台配置 U-Mamba 环境并训练医学图像分割

在 AutoDL 平台配置 U-Mamba 环境并训练医学图像分割

相关数据和离线依赖包，网盘下载链接
链接：https://pan.baidu.com/s/1DNjtsDJOlN_4l4Vr0y9tNg?pwd=bz95
提取码：bz95

视频版

环境配置教程来源
https://blog.csdn.net/weixin_45231460/article/details/138862816
U-Mamba项目地址
https://wanglab.ai/u-mamba.html

1.U-Mamba

最近，处理长序列任务的 Mamba 模型较为火爆。U-Mamba 的作者设计了一个混合的 CNN-SSM（卷积神经网络 - 状态空间序列模型）块，它结合了 CNN 的局部特征提取能力以及 SSM 捕捉长距离依赖的能力，还充分利用了 UNet 的 U 型结构的优势，从而提出了 U-Mamba 网络。并且，结合 nnUnetV2 框架，该网络能够自适应各种医学分割数据。

在腹部 CT 和腹部 MRI 数据上，其效果优于 nnUNet。作者还解释，U - Mamba 分割结果中的异常值更少。例如，在腹部 CT 扫描中能生成更准确的肝脏和胃的分割 mask；在 MRI 扫描中能更好地描绘胆囊的边界。同时，在面对异质外观时，它表现得更为稳健，比如在腹部 MRI 扫描中能更好地区分脾脏的外观。

2.在AutoDL配置环境

配置这个环境，一方面是想跑通U-Mamba，一方面U-Mamba是结合了nnUnetV2，即使不跑U-Mamba，也可以在本地没有足够显卡的情况下，在AutoDL平台跑通nnUnetV2来分割自己的数据。

配置环境，最烦就是安装Mamba的causal-conv1d和mamba_ssm，根据U-Mamba官网指示来安装会报错（我），根据这个博主来安装Mamba，可以顺利跑通U-Mamba

https://blog.csdn.net/weixin_45231460/article/details/138862816

3.数据

这次分割的数据是论文中展示的腹部MRI数据，分割13个器官

4.安装环境

1.登录autoDL，开一个3090的卡。

2.安装依赖包

#切换目录
cd autodl-tmp
#创建环境变量
conda create -n umamba python=3.10-y
#刷新环境变量
conda init bash && source /root/.bashrc
#激活环境
conda activate umamba 
#安装pytorch库
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1-i https://mirrors.aliyun.com/pypi/simple/#安装causal_conv1d 和 mamba，这里选择离线安装#下载离线包（已经下载）
wget https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.1.3.post1/causal_conv1d-1.1.3.post1+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
wget https://github.com/state-spaces/mamba/releases/download/v1.1.1/mamba_ssm-1.1.1+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
#上传并安装离线包
pip install causal_conv1d-1.1.3.post1+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl 
pip install mamba_ssm-1.1.1+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
#验证import torch
import mamba_ssm
#克隆U-Mamba项目
git clone https://github.com/bowang-lab/U-Mamba.git
#安装umamba 所需要的其他依赖包
cd U-Mamba/umamba 
pip install -e .

5.上传和处理数据

1.数据结构

data/
└─dataset.json        <u>#包含数据类别，数量和模态等信息</u>
└─nnUNet_preprocessed  <u>#存放预处理数据</u>
└─nnUNet_results              <u>#存放训练模型结果 </u>
└─nnUNet_raw                   <u>#存放原始数据</u>
    └─Dataset702_AbdomenMR   <u>#数据集命名，Dataset三位数_数据名(Dataset703_ABC)</u>
        ├─imagesTr                 #训练用的原始数据，每个名字补_0000
             ├─amos_0507_0000.nii.gz  
             ├─amos_0608_0000.nii.gz
        ├─imagesTs                  <u>#测试用的原始数据（非必需）</u>
            ├─amos_0901_0000.nii.gz  
            ├─amos_0902_0000.nii.gz
        ├─labelsTr                    <u>#训练用的原始数据对应的mask</u>
            ├─amos_0507.nii.gz  
            ├─amos_0608.nii.gz
        ├─labelsTs                    <u>#测试用的原始数据对应的mask（非必需）</u>
            ├─amos_0901.nii.gz  
            ├─amos_0902.nii.gz
        ├─predictTs                <u>#存放推理mask结果（非必需）</u>

2.上传data.zip到U-mamba目录下，并解压数据

cd .. 并 unzip -d ./data data.zip

3.配置环境变量（换成你的真实路径）。终端输入⬇️

export nnUNet_raw="autodl-tmp/U-Mamba/data/nnUNet_raw"
export nnUNet_preprocessed="autodl-tmp/U-Mamba/data/nnUNet_preprocessed"
export nnUNet_results="autodl-tmp/U-Mamba/data/nnUNet_results"

4.预处理数据

702是“Dataset702_AbdomenMR”中的702

nnUNetv2_plan_and_preprocess -d 702--verify_dataset_integrity

5.训练

all代表训练5折。0~4代表单独训练其中某一折

#使用UMamba模型训练
nnUNetv2_train 702 3d_fullres all-tr nnUNetTrainerUMambaBot
或
#不使用U-Mamba，用nnUNet 的 3D全分辨率U-Net训练
nnUNetv2_train 702 3d_fullres 0#训练过程保存的目录路径#autodl-tmp/U-Mamba/data/nnUNet_results/Dataset702_AbdomenMR/nnUNetTrainerUMambaBot__nnUNetPlans__3d_fullres/fold_all/checkpoint_best.pth

6.推理

-i 接数据输入目录

-o 接预测输出目录

-f 接推理用第几折，或者全用

-chk 接推理所需模型，默认checkpoint_final.pth（因为我没有完全训练完，没有生成checkpoint_final.pth，因此我指定读取checkpoint_best.pth）

#使用UMamba模型推理
nnUNetv2_predict -i ~/autodl-tmp/U-Mamba/data/nnUNet_raw/Dataset702_AbdomenMR/imagesTs -o ~/autodl-tmp/U-Mamba/data/nnUNet_raw/Dataset702_AbdomenMR/predictTs -d 702-c 3d_fullres -f all-tr nnUNetTrainerUMambaBot -chk checkpoint_best.pth
#使用nnUNet 的 3D全分辨率U-Net模型推理
nnUNetv2_predict -i ~/autodl-tmp/U-Mamba/data/nnUNet_raw/Dataset702_AbdomenMR/imagesTs -o ~/autodl-tmp/U-Mamba/data/nnUNet_raw/Dataset702_AbdomenMR/predictTs -d 702-c 3d_fullres -f 0-chk checkpoint_best.pth