0 前言0.1 UMambaUMamba是由Jun Ma、Feifei Li及Bo Wang于2024年提出的CNN-SSM混合分割架构论文《U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation》。其开创性地将mamba引入三维分割任务中缓解了传统CNN面临的感受野不足的局限性。框架在nnUNet的基础上开发充分发挥了nnUNet框架在数据处理策略等方面的优势。0.2 本教程面向人群UMamba基于nnUNet框架开发它继承了nnUNet众多特性但在网络实现方面更加简化和易于接入且性能波动不大参见《nnU-Net Revisited: A Call for Rigorous Validation in 3D Medical Image Segmentation》Table 1因此推荐希望基于mamba设计分割模型的研究者希望在nnUNet框架中更简单地部署分割模型的研究者使用。如果任务目标是快速落地一个可用的分割模型nnUNet仍然是更为快捷方便的选择。1 仓库1.1 仓库地址UMamba发布于github仓库U-Mamba中。推荐优先使用git clone https://github.com/bowang-lab/U-Mamba进行获取。1.2 仓库结构本节主要介绍UMamba仓库中需要关注的目录和文件。未介绍到的文件在初步接触框架时无需过多关注。1.2.1 数据集UMamba框架待处理的数据集统一位于U-Mamba/data中配置数据集的方法与nnUNet一致将在后文详细介绍。1.2.2 网络UMamba的网络实现位于U-Mamba/umamba/nnunetv2/nets中初始提供了UMambaBot_2d.py、UMambaBot_3d.py、UMambaEnc_2d.py和UMambaEnc_3d.py四种结构。其中UMambaBot仅在bottleneck上添加mambalayer显存占用更低UMambaEnc是论文的实现在编码器和bottleneck上均使用mamba layer。需要强调的是虽然论文中没有明确提及设计意图但UMambaEnc使用了空间mamba和通道mamba两种mamba当通道数≥空间token总数时mamba将沿通道扫描。do_channel_token [False] * n_stages feature_map_sizes [] feature_map_size input_size for s in range(n_stages): feature_map_sizes.append([i // j for i, j in zip(feature_map_size, strides[s])]) feature_map_size feature_map_sizes[-1] if np.prod(feature_map_size) features_per_stage[s]: do_channel_token[s] True1.2.3 训练器通过继承和重写训练器相关方法可将框架扩展到不同模型上。UMamba框架的训练器位于U-Mamba/umamba/nnunetv2/training/nnUNetTrainer内包括nnUNetTrainerUMambaBot.py、nnUNetTrainerUMambaEnc.py、nnUNetTrainerUMambaEncNoAMP.py。对于nnUNetTrainerUMambaEncNoAMP.py官方仓库给出解释如下AMP could lead to nan in the Mamba module. We also provide a trainer without AMP: U-Mamba/umamba/nnunetv2/training/nnUNetTrainer/nnUNetTrainerUMambaEncNoAMP.py at main · bowang-lab/U-Mamba · GitHub如果需要训练Bot模型而遇到相关问题可考虑仿照nnUNetTrainerUMambaEncNoAMP.py实现一个nnUNetTrainerUMambaBotNoAMP.py。2 环境配置本节基于Anaconda进行环境配置确保conda及pip使用了可用的源。环境配置并不唯一本节旨在给出一个可用的版本组合。2.1 操作系统mamba2只支持Linux环境若希望在Windows上使用UMamba框架建议使用WSL。本文使用Ubuntu 24.04.4 LTS进行演示。WSL的环境配置流程与Linux基本一致不再单独介绍。2.2 创建conda环境conda create -n umamba python3.10 -y conda activate umamba2.3 安装pytorchpytorch版本和支持的cuda版本需要根据自己的显卡确定例如官方仓库中给出的cu118在Blackwell架构50系显卡下一般不受支持需要cu128以上及对应支持版本的torch。本文安装torch2.1.1cu118版本在RTX 3080 Ti和RTX 4090 D上均可用。pip install torch2.1.1 torchvision0.16.1 --index-url https://download.pytorch.org/whl/cu1182.4 安装mamba不建议使用pip直接从pypi源安装mamba。获取github仓库所需的发布版mamba2https://github.com/state-spaces/mamba/releasescausal-conv1dhttps://github.com/Dao-AILab/causal-conv1d/releases本文采用的版本为:mambacausal-conv1d2.2.11.3.0.post1在仓库内发布版的命名一般形如mamba_ssm-2.2.1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl。其中cu118torch2.1字段表示其支持的torch版本需要与已安装的torch版本一致。cp310表示其支持的python版本需要与环境python版本一致。linux_x86_64表示系统及CPU架构需要与使用环境一致。cxx11abiFALSE字段表示该wsl文件是否使用GNU C11 ABI编译需要与torch的ABI一致通常从pypi源直接获取的torch的这项为False保险起见推荐通过以下代码检查torch的ABIimport torch print(torch._C._GLIBCXX_USE_CXX11_ABI)此处我们下载mamba_ssm-2.2.1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whlcausal_conv1d-1.3.0.post1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whlmamba在setup.py中指定了支持架构见line 176cc_flag.append(-gencode) cc_flag.append(archcompute_53,codesm_53) cc_flag.append(-gencode) cc_flag.append(archcompute_62,codesm_62) cc_flag.append(-gencode) cc_flag.append(archcompute_70,codesm_70) cc_flag.append(-gencode) cc_flag.append(archcompute_72,codesm_72) cc_flag.append(-gencode) cc_flag.append(archcompute_80,codesm_80) cc_flag.append(-gencode) cc_flag.append(archcompute_87,codesm_87) if bare_metal_version Version(11.8): cc_flag.append(-gencode) cc_flag.append(archcompute_90,codesm_90)这使得sm_60P100、sm_611080 Ti等架构的显卡无法得到支持。对于sm_61参见issue #40及issue #438的讨论hhhhpaaa提供了一个重编译的whl版本。注意按照要求安装triton-nightly 3.0.0.post20240626041721。如寻求更多架构的支持可尝试结合前述话题的内容自行编译。2.5 修复transformer版本mamba在setup.py中未显式指定transformers版本见line 370。这导致pip会优先安装transformers5.x的依赖对于本文安装的mamba2.2.1需要手动降级到4.x推荐版本为4.44.2。2.6 安装UMamba在1.2中我们介绍过UMamba框架的仓库结构它是在nnunetv2目录下进行模型和训练器维护的因此强烈建议以可编辑方式安装而非从pypi源获取nnunetv2。安装方法如下cd U-Mamba/umamba pip install -e .nnUNet在v2.4中引入了ResEnc配置而UMamba的标准运行版本是nnUNet v2.2。相比之下UMambaBot虽然使用残差卷积块但堆块策略与nnUNet ResEnc存在较大不同如果涉及实验需谨慎对待这部分差异。2.7 修复numpy版本这个步骤建议在最后完成。torch2.0、torch2.1这些边界版本可能无法正确约束numpy版本pip会错误地解析并安装numpy2需要手动将numpy版本回退到1.2x.x如1.26.3并处理某些提示的冲突如opencv-python。以下为带numpy版本约束的包安装示例pip install numpy1.26.3 opencv-python2.8 mamba3mamba3在2026年4月发布此处不展开介绍仅提供一组于RTX 4090 D上可运行的版本配置参考组合。库版本torch2.4.0cu124triton3.6.0mamba-ssm2.3.1causal-conv1d1.5.4transformers5.5.4numpy2.2.63 数据集制作本章介绍通过调用nnunetv2命令完成MSD数据集制作及配置规划。3.1 MSD转换如果你很确定你的数据集本身符合MSD标准则只需要创建一个DatasetXXX_DatasetName如Dataset001_brats并放入你的数据集无需进行以下步骤。在U-Mamba/data/nnUNet_raw目录下创建一个新文件夹DataXXX_DatasetName如Data001_brats。该文件夹内需要包含三个子目录和一个json文件。分别为imagesTr存储训练数据labelsTr存储训练数据标注训练数据和标注的名称需要一一对应imagesTs存储测试数据可为空文件夹因为测试不会自动进行不影响整个训练流程dataset.json数据集信息一个简单的json文件如下{ modality: { 0: CT }, labels: { 0: background, 1: target }, numTraining: 100, file_ending: .nii.gz, training: [], test:[] }框架默认的训练范式是n-fold交叉验证imagesTr的100例数据会被分为n折每折包括(n-1)*n/100个训练样本和n/100个验证样本各fold之间验证集不重叠这些验证样本用于监控训练移动平均伪diceissue #2128以获取最优ckpt会在每个fold训练结束时完整运行一次整图推理默认使用DSC指标。因此如果希望运行测试测试集需要在前述的imagesTs目录下单独预留并参与转换。如果只希望运行训练推理则该目录可以留空。上述步骤主要是为了让测试集参与转换测试集需要是MSD格式而位置并无强制性要求。事实上nnUNetv2 predict这一推理命令接受 -i 参数指定测试目录作为推理输入也接受 -d 参数指定数据集作为推理输入这是为了让一个已经训练完成的模型能够快速部署而无需持续维护Data目录。通过以下命令完成数据集转换nnUNetv2_convert_MSD_dataset -i ./DataXXX_DatasetName得到一个标准的MSD数据集目录DatasetXXX_DatasetName。其中json文件的内容为{ labels: { background: 0, target: 1 }, numTraining: 100, file_ending: .nii.gz, channel_names: { 0: CT } }这种MSD标准json文件无法被nnUNetv2_convert_MSD_dataset命令解析主要问题在于labels的键值对在转换前后映射是相反的因此不建议从Dataset701_AbdomenCT的参考json中直接构建json文件。此外channel_names键、modality键的语义一般被约定为模态名但对于单模态多通道数据可使用RGB代替模态。3.2 规划器和配置文件完成MSD数据集目录DatasetXXX_DatasetName的制作后需要通过nnunet规划器生成配置文件网络训练的众多参数如patchsize、batchsize、像素间距、堆块数在此决定。UMamba框架无需和nnUNet默认流程一样将数据集目录添加到环境变量它在path.py中硬编码了一组路径也即前述的U-Mamba/data/...。通过以下命令运行规划器nnUNetv2_plan_and_preprocess -d XXX --verify_dataset_integrity这个过程可能面临内存不足的情况对于Linux而言可尝试修改Swap或指定-np参数。np参数代表规划所使用进程数2D和3D lowres默认为83D fullres和其它配置默认为4修改到2或1能够缓解内存占用过高的问题代价是规划速度变慢。规划完成后将在U-Mamba/nnUNet_preprocessed目录下找到DatasetXXX_DatasetName子目录。其中nnUNetPlans.json为配置文件。configurations键包含了全部网络类型默认为2d、3d_fullres当级联启用时还存在3d_lowres配置。每个配置中需要重点关注的配置项包括batchsize批次大小不论显存情况如何这项数值不推荐低于2否则可能影响统计稳定性。patchsize分块大小nnUNet的训练基于patchpatch过小可能导致全局覆盖率不足patch过大可能导致显存占用上升。spacing像素/体素间距值越大轴向分辨率越低若修改patchsize建议对spacing进行同步修改。n_conv_per_stage_encoder表示网络的层数stage以及编码器每层卷积块的数量在标准的nnUNet v2.2中此处增加stage或卷积块数量可能引发OOM。但由于UMamba并未对mamba layer使用规划器评估显存占用因此框架本身就存在OOM的风险。综上建议此处可以根据实验需求进行修改但需要做好显存管理措施。n_conv_per_stage_decoder表示网络解码器每层卷积块的数量同上。此外UMamba限制了编解码器深层卷积块的数量恒定为1line 437和line 396这是缓解OOM的静态措施如有需要可进行调整。4 训练在完成数据集制作后若数据集为DatasetXXX_DatasetName可通过以下命令启动训练nnUNetv2_train XXX 3d_fullres 0 -tr nnUNetTrainerUMambaEnc #e.g. train fold 0 with 3d_fullres config, using UMambaEnc nnUNetv2_train XXX 2d 1 -tr nnUNetTrainerUMambaBot #e.g. train fold 1 with 2d config, using UMambaBot通常当epoch 0训练结束输出如图所示的结果时说明模型已经能够正常进行训练了。部分超参数如epoch数量、学习率等可在nnUNetTrainer.py line 142-149设置。若出现OOM报错可尝试以下解决方案减少batchsize至不低于2或采用梯度累积减少patchsize并调整spacing优先降采样高分辨率轴训练预热以应对可能存在的初始峰值使用checkpoint对模块进行封装训练时长上升启用AMP可能与mamba冲突禁用深监督可能影响训练效果5 UMamba框架与nnUNet框架区别总结UMamba通过nets目录直接维护网络可扩展性较高UMamba在深层减少堆块nnUNet在v2.4鼓励在深层堆块UMamba在默认配置下可能发生OOMnnUNet一般不会相比nnUNet v2.2UMamba默认使用残差块UMamba受限于mamba依赖无法在Windows下直接运行UMamba重构了数据集路径解析策略使得不同可编辑库独立可同时运行不会在环境变量上产生竞争UMamba环境配置较为复杂参考资料GitHub仓库GitHub - MIC-DKFZ/nnUNet · GitHubGitHub - bowang-lab/U-Mamba: U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation · GitHubGitHub - state-spaces/mamba: Mamba SSM architecture · GitHubGitHub - Dao-AILab/causal-conv1d: Causal depthwise conv1d in CUDA, with a PyTorch interface · GitHubGitHub - hhhhpaaa/Mamba-ssm: A well-built wheel is used to support the Nvidia 10 series graphics card. · GitHub论文Ma J, Li F, Wang B. U-mamba: Enhancing long-range dependency for biomedical image segmentation[J]. arXiv preprint arXiv:2401.04722, 2024.Isensee F, Wald T, Ulrich C, et al. nnu-net revisited: A call for rigorous validation in 3d medical image segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024: 488-498. UMamba部署及训练教程
0 前言0.1 UMambaUMamba是由Jun Ma、Feifei Li及Bo Wang于2024年提出的CNN-SSM混合分割架构论文《U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation》。其开创性地将mamba引入三维分割任务中缓解了传统CNN面临的感受野不足的局限性。框架在nnUNet的基础上开发充分发挥了nnUNet框架在数据处理策略等方面的优势。0.2 本教程面向人群UMamba基于nnUNet框架开发它继承了nnUNet众多特性但在网络实现方面更加简化和易于接入且性能波动不大参见《nnU-Net Revisited: A Call for Rigorous Validation in 3D Medical Image Segmentation》Table 1因此推荐希望基于mamba设计分割模型的研究者希望在nnUNet框架中更简单地部署分割模型的研究者使用。如果任务目标是快速落地一个可用的分割模型nnUNet仍然是更为快捷方便的选择。1 仓库1.1 仓库地址UMamba发布于github仓库U-Mamba中。推荐优先使用git clone https://github.com/bowang-lab/U-Mamba进行获取。1.2 仓库结构本节主要介绍UMamba仓库中需要关注的目录和文件。未介绍到的文件在初步接触框架时无需过多关注。1.2.1 数据集UMamba框架待处理的数据集统一位于U-Mamba/data中配置数据集的方法与nnUNet一致将在后文详细介绍。1.2.2 网络UMamba的网络实现位于U-Mamba/umamba/nnunetv2/nets中初始提供了UMambaBot_2d.py、UMambaBot_3d.py、UMambaEnc_2d.py和UMambaEnc_3d.py四种结构。其中UMambaBot仅在bottleneck上添加mambalayer显存占用更低UMambaEnc是论文的实现在编码器和bottleneck上均使用mamba layer。需要强调的是虽然论文中没有明确提及设计意图但UMambaEnc使用了空间mamba和通道mamba两种mamba当通道数≥空间token总数时mamba将沿通道扫描。do_channel_token [False] * n_stages feature_map_sizes [] feature_map_size input_size for s in range(n_stages): feature_map_sizes.append([i // j for i, j in zip(feature_map_size, strides[s])]) feature_map_size feature_map_sizes[-1] if np.prod(feature_map_size) features_per_stage[s]: do_channel_token[s] True1.2.3 训练器通过继承和重写训练器相关方法可将框架扩展到不同模型上。UMamba框架的训练器位于U-Mamba/umamba/nnunetv2/training/nnUNetTrainer内包括nnUNetTrainerUMambaBot.py、nnUNetTrainerUMambaEnc.py、nnUNetTrainerUMambaEncNoAMP.py。对于nnUNetTrainerUMambaEncNoAMP.py官方仓库给出解释如下AMP could lead to nan in the Mamba module. We also provide a trainer without AMP: U-Mamba/umamba/nnunetv2/training/nnUNetTrainer/nnUNetTrainerUMambaEncNoAMP.py at main · bowang-lab/U-Mamba · GitHub如果需要训练Bot模型而遇到相关问题可考虑仿照nnUNetTrainerUMambaEncNoAMP.py实现一个nnUNetTrainerUMambaBotNoAMP.py。2 环境配置本节基于Anaconda进行环境配置确保conda及pip使用了可用的源。环境配置并不唯一本节旨在给出一个可用的版本组合。2.1 操作系统mamba2只支持Linux环境若希望在Windows上使用UMamba框架建议使用WSL。本文使用Ubuntu 24.04.4 LTS进行演示。WSL的环境配置流程与Linux基本一致不再单独介绍。2.2 创建conda环境conda create -n umamba python3.10 -y conda activate umamba2.3 安装pytorchpytorch版本和支持的cuda版本需要根据自己的显卡确定例如官方仓库中给出的cu118在Blackwell架构50系显卡下一般不受支持需要cu128以上及对应支持版本的torch。本文安装torch2.1.1cu118版本在RTX 3080 Ti和RTX 4090 D上均可用。pip install torch2.1.1 torchvision0.16.1 --index-url https://download.pytorch.org/whl/cu1182.4 安装mamba不建议使用pip直接从pypi源安装mamba。获取github仓库所需的发布版mamba2https://github.com/state-spaces/mamba/releasescausal-conv1dhttps://github.com/Dao-AILab/causal-conv1d/releases本文采用的版本为:mambacausal-conv1d2.2.11.3.0.post1在仓库内发布版的命名一般形如mamba_ssm-2.2.1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl。其中cu118torch2.1字段表示其支持的torch版本需要与已安装的torch版本一致。cp310表示其支持的python版本需要与环境python版本一致。linux_x86_64表示系统及CPU架构需要与使用环境一致。cxx11abiFALSE字段表示该wsl文件是否使用GNU C11 ABI编译需要与torch的ABI一致通常从pypi源直接获取的torch的这项为False保险起见推荐通过以下代码检查torch的ABIimport torch print(torch._C._GLIBCXX_USE_CXX11_ABI)此处我们下载mamba_ssm-2.2.1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whlcausal_conv1d-1.3.0.post1cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whlmamba在setup.py中指定了支持架构见line 176cc_flag.append(-gencode) cc_flag.append(archcompute_53,codesm_53) cc_flag.append(-gencode) cc_flag.append(archcompute_62,codesm_62) cc_flag.append(-gencode) cc_flag.append(archcompute_70,codesm_70) cc_flag.append(-gencode) cc_flag.append(archcompute_72,codesm_72) cc_flag.append(-gencode) cc_flag.append(archcompute_80,codesm_80) cc_flag.append(-gencode) cc_flag.append(archcompute_87,codesm_87) if bare_metal_version Version(11.8): cc_flag.append(-gencode) cc_flag.append(archcompute_90,codesm_90)这使得sm_60P100、sm_611080 Ti等架构的显卡无法得到支持。对于sm_61参见issue #40及issue #438的讨论hhhhpaaa提供了一个重编译的whl版本。注意按照要求安装triton-nightly 3.0.0.post20240626041721。如寻求更多架构的支持可尝试结合前述话题的内容自行编译。2.5 修复transformer版本mamba在setup.py中未显式指定transformers版本见line 370。这导致pip会优先安装transformers5.x的依赖对于本文安装的mamba2.2.1需要手动降级到4.x推荐版本为4.44.2。2.6 安装UMamba在1.2中我们介绍过UMamba框架的仓库结构它是在nnunetv2目录下进行模型和训练器维护的因此强烈建议以可编辑方式安装而非从pypi源获取nnunetv2。安装方法如下cd U-Mamba/umamba pip install -e .nnUNet在v2.4中引入了ResEnc配置而UMamba的标准运行版本是nnUNet v2.2。相比之下UMambaBot虽然使用残差卷积块但堆块策略与nnUNet ResEnc存在较大不同如果涉及实验需谨慎对待这部分差异。2.7 修复numpy版本这个步骤建议在最后完成。torch2.0、torch2.1这些边界版本可能无法正确约束numpy版本pip会错误地解析并安装numpy2需要手动将numpy版本回退到1.2x.x如1.26.3并处理某些提示的冲突如opencv-python。以下为带numpy版本约束的包安装示例pip install numpy1.26.3 opencv-python2.8 mamba3mamba3在2026年4月发布此处不展开介绍仅提供一组于RTX 4090 D上可运行的版本配置参考组合。库版本torch2.4.0cu124triton3.6.0mamba-ssm2.3.1causal-conv1d1.5.4transformers5.5.4numpy2.2.63 数据集制作本章介绍通过调用nnunetv2命令完成MSD数据集制作及配置规划。3.1 MSD转换如果你很确定你的数据集本身符合MSD标准则只需要创建一个DatasetXXX_DatasetName如Dataset001_brats并放入你的数据集无需进行以下步骤。在U-Mamba/data/nnUNet_raw目录下创建一个新文件夹DataXXX_DatasetName如Data001_brats。该文件夹内需要包含三个子目录和一个json文件。分别为imagesTr存储训练数据labelsTr存储训练数据标注训练数据和标注的名称需要一一对应imagesTs存储测试数据可为空文件夹因为测试不会自动进行不影响整个训练流程dataset.json数据集信息一个简单的json文件如下{ modality: { 0: CT }, labels: { 0: background, 1: target }, numTraining: 100, file_ending: .nii.gz, training: [], test:[] }框架默认的训练范式是n-fold交叉验证imagesTr的100例数据会被分为n折每折包括(n-1)*n/100个训练样本和n/100个验证样本各fold之间验证集不重叠这些验证样本用于监控训练移动平均伪diceissue #2128以获取最优ckpt会在每个fold训练结束时完整运行一次整图推理默认使用DSC指标。因此如果希望运行测试测试集需要在前述的imagesTs目录下单独预留并参与转换。如果只希望运行训练推理则该目录可以留空。上述步骤主要是为了让测试集参与转换测试集需要是MSD格式而位置并无强制性要求。事实上nnUNetv2 predict这一推理命令接受 -i 参数指定测试目录作为推理输入也接受 -d 参数指定数据集作为推理输入这是为了让一个已经训练完成的模型能够快速部署而无需持续维护Data目录。通过以下命令完成数据集转换nnUNetv2_convert_MSD_dataset -i ./DataXXX_DatasetName得到一个标准的MSD数据集目录DatasetXXX_DatasetName。其中json文件的内容为{ labels: { background: 0, target: 1 }, numTraining: 100, file_ending: .nii.gz, channel_names: { 0: CT } }这种MSD标准json文件无法被nnUNetv2_convert_MSD_dataset命令解析主要问题在于labels的键值对在转换前后映射是相反的因此不建议从Dataset701_AbdomenCT的参考json中直接构建json文件。此外channel_names键、modality键的语义一般被约定为模态名但对于单模态多通道数据可使用RGB代替模态。3.2 规划器和配置文件完成MSD数据集目录DatasetXXX_DatasetName的制作后需要通过nnunet规划器生成配置文件网络训练的众多参数如patchsize、batchsize、像素间距、堆块数在此决定。UMamba框架无需和nnUNet默认流程一样将数据集目录添加到环境变量它在path.py中硬编码了一组路径也即前述的U-Mamba/data/...。通过以下命令运行规划器nnUNetv2_plan_and_preprocess -d XXX --verify_dataset_integrity这个过程可能面临内存不足的情况对于Linux而言可尝试修改Swap或指定-np参数。np参数代表规划所使用进程数2D和3D lowres默认为83D fullres和其它配置默认为4修改到2或1能够缓解内存占用过高的问题代价是规划速度变慢。规划完成后将在U-Mamba/nnUNet_preprocessed目录下找到DatasetXXX_DatasetName子目录。其中nnUNetPlans.json为配置文件。configurations键包含了全部网络类型默认为2d、3d_fullres当级联启用时还存在3d_lowres配置。每个配置中需要重点关注的配置项包括batchsize批次大小不论显存情况如何这项数值不推荐低于2否则可能影响统计稳定性。patchsize分块大小nnUNet的训练基于patchpatch过小可能导致全局覆盖率不足patch过大可能导致显存占用上升。spacing像素/体素间距值越大轴向分辨率越低若修改patchsize建议对spacing进行同步修改。n_conv_per_stage_encoder表示网络的层数stage以及编码器每层卷积块的数量在标准的nnUNet v2.2中此处增加stage或卷积块数量可能引发OOM。但由于UMamba并未对mamba layer使用规划器评估显存占用因此框架本身就存在OOM的风险。综上建议此处可以根据实验需求进行修改但需要做好显存管理措施。n_conv_per_stage_decoder表示网络解码器每层卷积块的数量同上。此外UMamba限制了编解码器深层卷积块的数量恒定为1line 437和line 396这是缓解OOM的静态措施如有需要可进行调整。4 训练在完成数据集制作后若数据集为DatasetXXX_DatasetName可通过以下命令启动训练nnUNetv2_train XXX 3d_fullres 0 -tr nnUNetTrainerUMambaEnc #e.g. train fold 0 with 3d_fullres config, using UMambaEnc nnUNetv2_train XXX 2d 1 -tr nnUNetTrainerUMambaBot #e.g. train fold 1 with 2d config, using UMambaBot通常当epoch 0训练结束输出如图所示的结果时说明模型已经能够正常进行训练了。部分超参数如epoch数量、学习率等可在nnUNetTrainer.py line 142-149设置。若出现OOM报错可尝试以下解决方案减少batchsize至不低于2或采用梯度累积减少patchsize并调整spacing优先降采样高分辨率轴训练预热以应对可能存在的初始峰值使用checkpoint对模块进行封装训练时长上升启用AMP可能与mamba冲突禁用深监督可能影响训练效果5 UMamba框架与nnUNet框架区别总结UMamba通过nets目录直接维护网络可扩展性较高UMamba在深层减少堆块nnUNet在v2.4鼓励在深层堆块UMamba在默认配置下可能发生OOMnnUNet一般不会相比nnUNet v2.2UMamba默认使用残差块UMamba受限于mamba依赖无法在Windows下直接运行UMamba重构了数据集路径解析策略使得不同可编辑库独立可同时运行不会在环境变量上产生竞争UMamba环境配置较为复杂参考资料GitHub仓库GitHub - MIC-DKFZ/nnUNet · GitHubGitHub - bowang-lab/U-Mamba: U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation · GitHubGitHub - state-spaces/mamba: Mamba SSM architecture · GitHubGitHub - Dao-AILab/causal-conv1d: Causal depthwise conv1d in CUDA, with a PyTorch interface · GitHubGitHub - hhhhpaaa/Mamba-ssm: A well-built wheel is used to support the Nvidia 10 series graphics card. · GitHub论文Ma J, Li F, Wang B. U-mamba: Enhancing long-range dependency for biomedical image segmentation[J]. arXiv preprint arXiv:2401.04722, 2024.Isensee F, Wald T, Ulrich C, et al. nnu-net revisited: A call for rigorous validation in 3d medical image segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024: 488-498.