在 AutoDL 平台配置 U-Mamba 环境并训练自定义2d数据集

在 AutoDL 平台配置 U-Mamba 环境并训练自定义2d数据集

需要在 AutoDL 平台配置 U-Mamba 环境并训练自定义2d数据集，看了一些博客的复现，为了方便一直训练，所以在此记录一下

1 环境安装

###1.1 pytorch 与 nnunet
环境安装可看此博客
###1.2 autodl配置
在autodl服务器上配置
###1.3 autodl与pycharm联通
autodl与pycharm联通

配置自定义数据集

###2.1 数据集格式

• 数据集dataset_name的格式： DatasetID_名称（ 数字大于10）为数据集的ID
• 相比于labels数据集，images数据集中原图命名后多加_0000，png格式

###2.2 创建需要的目录

• eval_result 评估结果
• nnUNet_predict_result 推理结果
• nnUNet_preprocessed 数据预处理之后的图像
• nnUNet_raw 原始数据集
• nnUNet_results 训练模型即验证结果

数据预处理

###3.1 配置环境变量(换成你的真实路径）

首先

vim .bashrc

添加以下三行

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

保存 Esc

:wq

最后

source .bashrc

###3.2 数据预处理

nnUNetv2_plan_and_preprocess -d 110 --verify_dataset_integrit

DATASET_ID（即为自选数据集ID，注意原图和标签的ID要一一对应）

训练

nnUNetv2_train 110 2d 9 -tr nnUNetTrainerUMambaEAP

• 110即为自定义数据集ID
• 9代表十折交叉验证（可更改验证方式，在模型主函数里修改do_split函数后重新划分数据集即可，默认是五折交叉）
• nnUNetTrainerUMambaEAP为自定义网络模型，也可替换成别的网络模型

推理

###4.1 寻找最佳配置
在nnunetv2/evaluation/find_best_configuration.py 更改主函数中的数据集ID为你的自定义ID，交叉验证改为你选择的折数，然后运行

 cd /U-Mamba-main/umamba/nnunetv2/evaluation
 python find_best_configuration.py

###4.2 推理

nnUNetv2_predict -i ~/autodl-tmp/U-Mamba-main/data/nnUNet_raw/Dataset110_OCTASegmentation/imagesTs -o ~/autodl-tmp/U-Mamba-main/data/nnUNet_predict_result/Dataset110_result -d 110 -c 2d -f 9 -tr nnUNetTrainerUMambaEAP --disable_tta -chk checkpoint_best.pth

• -i 原图存储的位置
• -o 推理结果存储的位置；
• 2d 数据集2d类型：
• -f 接推理用第几折，或者全用；
• -chk 预测模型，默认为checkpint_final.pth

评估

借用nnUNetV2的源码文件进行评估，位置：nnunetv2/evaluation/evaluate_predictions.py
对主函数进行修改， 修改完成后运行即可，得到DICE与Iou两个指标

• floder_ref为测试集GT文件；floder_pred 为推理结果数据集；output_file 为需要输出的评价结果，json文件
• labels_to_list_of_regions([1]) 前景是1，指定需要评估的区域为label1
• 文件后缀，修改为.png文件

folder_ref = '/root/autodl-tmp/U-Mamba-main/data/nnUNet_raw/Dataset110_OCTASegmentation/labelsTs'
folder_pred = '/root/autodl-tmp/U-Mamba-main/data/nnUNet_predict_result/Dataset110_result'
output_file = '/root/autodl-tmp/U-Mamba-main/data/eval_result/topo.json'

image_reader_writer = SimpleITKIO()
file_ending = '.png'
regions = labels_to_list_of_regions([1])
ignore_label = None
num_processes = 12