如果你最近在复现医学图像分割方向的论文，大概率绕不开 U-Mamba 这个名字。它是在 nnU-Net 基础上扩展出来的 3D 分割框架，算是“代码复现”圈子里热度很高的一份工作。作为一个跑过 U-Mamba 完整训练流程的人，我可以很负责任地说：整个复现过程里最容易让人崩溃的，不是模型结构，不是训练策略，而是第一步“创建环境”。很多人三天没跑通代码，最后发现只是 conda 环境里 Python 版本不对，或者 CUDA 和 PyTorch 对不上号。这篇就专门聊透 Umamba 代码复现中的创建环境环节，把环境清单、命令写法、依赖安装顺序和坑点排查一次讲清楚。适合正在复现论文的研究生、刚接触 3D 医学图像分割的开发者，以及所有被环境问题折磨过的朋友。
1. 复现U-Mamba之前，先搞清楚它要什么样的环境
1.1 U-Mamba的定位决定了环境清单
U-Mamba 不是那种“从零手写”的独立模型，它本质上是基于 nnU-Net（Medical Segmentation Decathlon 时期就积累下来的 3D 分割框架）改造而来。核心思想是在 U-Net 的编码器和解码器之间引入状态空间模型（SSM，也就是 Mamba），用来增强长距离依赖建模能力。这个定位决定了环境搭建不是装一个库那么简单，而是一整套配套生态：Python 解释器、PyTorch、MONAI、nnU-Net、U-Mamba 源码包，以及底层的 CUDA 驱动和 GPU 硬件。
很多人第一次看到 U-Mamba 的 GitHub 仓库，会下意识以为直接pip install umamba就行了，真这样做大概率会碰壁。因为它的安装脚本会依赖 nnU-Net，而 nnU-Net v2 的安装又涉及pip install -e这种可编辑模式，还会要求你设置nnUNet_raw、nnUNet_preprocessed、nnUNet_results三个环境变量。U-Mamba 的复现环境本质上等于“nnU-Net 环境 + 一堆额外的依赖”，官方 README 里的安装顺序也是先装 nnU-Net，再装 U-Mamba。所以环境清单里至少要有：
- Python 3.9 或 3.10
- PyTorch 2.x（带 CUDA 支持）
- MONAI（U-Mamba 数据处理和增强会用到）
- nnU-Net v2 源码（以可编辑模式安装）
- U-Mamba 源码（同样以可编辑模式安装）
- NVIDIA GPU 驱动 + CUDA 工具链
搞清楚这个定位，后面的每一步才不会抓瞎。这也是我想强调的：复现论文之前，先花十分钟读一遍仓库 README，看它依赖了哪些上游项目，比急着敲命令有用得多。
1.2 系统、GPU与显存：不同工作流的不同选择
环境这件事，首先要面对的是操作系统。U-Mamba 官方的训练和推理脚本主要面向 Linux 开发环境，Ubuntu 20.04 或 22.04 是我见过最顺的组合。Windows 不是完全不行，但你大概率会踩到编译不通过、nnUNetv2_train命令识别不了、路径分隔符风格不一致等一系列问题。我建议 Windows 用户老老实实启用 WSL2（Windows Subsystem for Linux）再装环境，体验会稳定很多。
GPU 方面，U-Mamba 是 3D 医学图像分割模型，处理的是三维体积数据，显存需求比普通 2D 模型大得多。我本地测试时用的一张 RTX 3090（24GB 显存）跑小 patch 训练还算流畅，但如果直接上全分辨率原始数据，24GB 也可能不够。显存低于 8GB 的话，别指望完整复现论文里的训练配置，至少要把patch_size、batch_size降下来。这部分不是环境搭建本身，但它直接影响你创建环境时的 CUDA 版本选型：显存越大，越可以考虑新版本的 CUDA 和 PyTorch；老卡反而要谨慎，比如 20 系显卡对最新 CUDA 的支持就没那么完美。
顺带说一句，如果你只想先验证 U-Mamba 能不能跑通，不一定要追求和论文完全一致的硬件。我认识不少同学是先在云服务器上用它自带的demo脚本测一遍，再回本地做完整训练。环境搭建的原则是一致的：先追求“能跑”，再追求“跑得快”。
如果服务器或本地机器上已经装好了 Docker，也可以考虑直接用官方镜像，但镜像同样依赖宿主机有可用的 NVIDIA 驱动。所以我个人还是推荐先用 conda 把环境理清楚，这是后续所有容器化方案的基础。
2. 创建conda环境：版本选型与最稳命令
2.1 为什么环境管理我坚持用conda
网上现在有很多环境管理工具，比如 Python 自带的venv、Facebook 的uv，但我给复现论文的流程推荐conda，原因很简单：它不只是管 Python 包，还能管 CUDA 相关依赖、以及一些非 Python 的系统库。U-Mamba 这种项目，依赖的是 PyTorch、MONAI、SimpleITK、batchgenerators 这些包，彼此之间存在版本关联。如果全塞到系统 Python 里，迟早会出现“A 库要求 numpy 1.x，B 库已经用了 numpy 2.x”的冲突。
我用一个生活化的类比：conda 环境就像给每个项目准备一间独立厨房。你在“umamba”这间厨房里炒菜，锅碗瓢盆（依赖库）随便摆，不会影响旁边“pytorch_old”那间厨房里的厨具。哪怕你把 umamba 环境玩坏了，直接conda remove -n umamba --all删掉，换一个重来就行，系统里其他项目毫发无损。
还要提一个新手容易混淆的点：U-Mamba 框架里的“Mamba”是状态空间模型，而环境管理工具“mamba”是 conda 的 C++ 重写加速版，两者只是名字撞车。你在 README 里看到conda create的命令，不一定非要装 mambaforge，用 miniconda 也完全可以。只是如果你发现conda create卡在 “Solving environment” 这一步很久，可以考虑装一个 mamba 来提速，后面我会专门讲这个坑。
2.2 Python版本到底选3.9还是3.10，我实测后的结论
U-Mamba 官方 README 的建议是 Python 3.9，而我实际测试下来，3.10 也能正常跑完依赖安装，只是要小心 PyTorch 和 MONAI 的版本上限。PyTorch 从 1.13 开始就完整支持 3.10，MONAI 早期版本对 3.10 的支持也基本没问题，所以理论上两个都能用。
但我还是推荐按官方文档走，选 Python 3.9。原因有两个：第一，nnU-Net v2 内部一些脚本依赖的SimpleITK和batchgenerators在 Python 3.9 下测试最充分，遇到奇怪报错的概率最低；第二，和论文复现的“可重复性”原则有关——官方所有测试是在 3.9 上做的，你如果为了图新鲜上 3.12，很容易踩到某个依赖还没适配新解释器的坑，最后还得回退。
如果你已经在电脑上装好了 Anaconda 或 Miniconda，创建环境的命令可以直接指定 Python 版本。如果还没装，先去 Miniconda 官网下载适配你系统的安装包。装完之后在终端执行：
conda create -n umamba python=3.9 -y conda activate umamba python --version看到Python 3.9.x的输出，就算第一步完成。这里有个小细节：-n umamba是环境名，可以自己改，但建议和你后续的代码目录、日志目录保持一致，这样不容易搞混。我自己就吃过亏，环境名叫test01，项目跑了一半才发现环境变量和数据集路径全都乱了。
2.3 可以直接抄作业的创建命令
很多人创建环境会直接在后面追加一大堆依赖，比如python=3.9 torch monai nnunet，我不建议这么做。原因是 conda 在解析这种大而全的依赖组合时非常慢，而且它默认从官方源拉包，国内网络环境下载速度一言难尽。更稳的做法是分三步走：先创建干净的基础环境，然后切入 PyTorch 安装，最后再装业务依赖。
第一步可以先配置 conda 的国内镜像源，如果你在国内环境下操作，这一步能省非常多时间。我一般在~/.condarc里写：
channels: - conda-forge - defaults show_channel_urls: true然后把conda config --set remote_read_timeout_secs 600也顺便加上。网络上偶尔会抽风，改大超时时限能少报几次错。接着再执行创建命令：
conda create -n umamba python=3.9 -y conda activate umamba这样生产出来的环境是干净的，后面你手动安装核心依赖时，版本搭配的主动权就完全在自己手里了。如果你真的想一条命令到底，也可以写conda create -n umamba python=3.9 pip -y，其实多塞一个 pip 进去也无妨，因为后面很多包还是得用 pip 装。
我特别想提醒一句：创建环境时千万别贪多，把 PyTorch 和 MONAI 直接写进conda create命令里。conda 默认的 PyTorch 通道版本更新滞后，很容易给你装上一个 CPU 版本或非常旧的 CUDA 版，等训练时报错就晚了。
3. 装依赖：PyTorch、nnU-Net、U-Mamba怎么装不打架
3.1 装PyTorch前必须摸清的三个问题
PyTorch 是整个环境里最核心又最容易翻车的部分。装它之前，先回答三个问题：你的显卡驱动支持到哪个 CUDA 版本？你需要的是 CPU 版还是 GPU 版？你后续用到的 MONAI、nnU-Net 和它兼容吗？
第一个问题用命令nvidia-smi就能看。输出里第一行有CUDA Version: 12.1这样的信息，这指的是驱动支持的最高 CUDA 版本，不是你已经装了 CUDA 工具包。PyTorch 的预编译包内置了对应版本的 CUDA 运行库，所以你不一定需要单独装完整的 CUDA Toolkit，只要驱动版本不低于 PyTorch 包要求的版本就行。比如你想用cu118版本的 PyTorch，那驱动至少得支持 CUDA 11.8；你想用cu121，驱动就得对应 12.1。
我现在的踩坑经验是：驱动如果是 535 系列或更新的，直接选cu121对应版本的 PyTorch；驱动是 470 系列的老卡，就老实选cu118。不建议一上来就追最新版 PyTorch，因为 nnU-Net v2 和 MONAI 对这种“尖端版本”的兼容性往往滞后。实测下来，PyTorch 2.x 任意一个稳定版本都可以，但没必要用 dev 版本。
安装命令我习惯用 PyTorch 官方生成的，以 CUDA 12.1 为例：
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果你在国内，PyTorch 官方源的下载速度会比较慢，可以先用国内镜像把包下载下来，或者把 PyTorch 的 wheel 文件提前下载到本地再pip install 本地文件路径。安装完后立刻验证：
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"如果输出2.x.x True，说明 GPU 版本已经就绪。这里最经典的坑是torch.cuda.is_available()返回False，很可能就是驱动版本太低，或者装的 torch 是 CPU 版。遇到这个问题先别急着重装，用python -c "import torch; print(torch.version.cuda)"看 torch 内置的 CUDA 版本是不是你预期的那一个。
3.2 从nnU-Net到U-Mamba的安装顺序与操作
U-Mamba 依赖 nnU-Net，所以安装顺序一定不能乱：先 nnU-Net，后 U-Mamba。顺序反了，U-Mamba 导入时找不到 nnU-Net 模块，报错会非常难懂。
推荐在 conda 环境里创建一个专门存放源码的目录，比如~/repos，然后按顺序执行：
cd ~/repos git clone https://github.com/MIC-DKFZ/nnUNet.git cd nnUNet pip install -e . cd ~/repos git clone https://github.com/himazin915/U-Mamba.git cd U-Mamba pip install -e .顺带一提，如果 GitHub 克隆速度很慢，可以直接在浏览器里下载 zip 包，解压后再pip install -e .，效果完全一样。这里用-e可编辑模式安装，好处是源码头文件一改就生效，后续你如果想调试 U-Mamba 内部代码，直接修改仓库里的.py文件即可，不用重装。
装完之后，nnU-Net v2 需要配置三个环境变量，不加的话，训练脚本找不到原始数据和预处理输出目录。我通常在~/.bashrc里加这几行：
export nnUNet_raw="/data/nnUNet_raw" export nnUNet_preprocessed="/data/nnUNet_preprocessed" export nnUNet_results="/data/nnUNet_results"然后把/data这些目录提前建好，再source ~/.bashrc。这一步很多复现教程一笔带过，但它其实是复现成功率最高的关键点之一。目录名不是固定的，你可以改成任何有读权限的位置，但三个路径不能是同一个，否则数据会被互相覆盖。
如果你在 Windows 的 WSL2 环境里操作，建议把数据放到 Linux 文件系统里，也就是/home下，别放到/mnt/c这种 Windows 挂载盘。挂在 NTFS 上的文件读写性能很差，训练时频繁 IO 会明显拖慢速度。
3.3 三步验证你的环境能不能跑
依赖装完不等于环境能用，我建议按顺序做三个验证动作：
第一，验证包是否都能导入。在终端输入：
python -c "import torch, monai, nnunetv2; print('base ok')" python -c "import umamba; print('umamba ok')"如果import umamba没有报错，说明 U-Mamba 代码已经装进当前环境并被解释器正确识别。
第二，验证 CUDA 真正可用。前面说过torch.cuda.is_available()要返回True，更严格的验证是用一个小张量做一次 GPU 运算：
import torch x = torch.rand(2, 2).cuda() print(x)能看到二维张量输出，说明 CUDA 通路正常。这里建议别只测is_available()，因为有些驱动异常会返回 True，但真正做乘法时会报错。
第三，验证 nnU-Net 的命令行工具能启动。运行：
nnUNetv2_plan_and_preprocess -h或者更简单一点：
which nnUNetv2_train能找到命令文件，基本说明 nnU-Net 安装成功。到这一步，U-Mamba 环境已经具备跑通 demo 的条件了。我见过不少人装完依赖就急着启动大训练，结果折腾一下午才发现 MONAI 版本太低导致数据增强出问题。所以这三步验证别偷懒。
另外提醒一句：之后在 PyCharm 或 VS Code 里运行时，记得把解释器换成umamba环境。PyCharm 可以在设置里搜“Python Interpreter”，选择~/miniconda3/envs/umamba/bin/python；VS Code 则在底部命令面板里选“Python: Select Interpreter”。如果 IDE 一直用的是系统解释器，你 conda 环境里装好的包一个都 import 不到，只会报红。
4. 实操中反复出现的坑与排查套路
4.1 下载慢和conda卡死
复现 U-Mamba 环境时最常见的问题不是代码逻辑，而是“网络不好、下载卡住”。conda create明明在转圈，却迟迟不结束，大部分时候是依赖解析太慢，尤其是 conda 官方源。解决办法有两个方向：一是给 conda 换成conda-forge或国内的镜像源，另一个是直接用 mamba 替代 conda 来创建环境。
mamba 的安装很轻量，装 Miniconda 之后可以这样装：
conda install -n base -c conda-forge mamba -y mamba create -n umamba python=3.9 -ymamba create在解析依赖时用的是 C++ 实现的求解器，速度比原生 conda 快很多。我也尝试过直接用uv来管理环境，但考虑 U-Mamba 依赖的包比较传统，最后大部分时候还是用 mamba 更稳。
如果卡在下载而不是解析阶段，可以检查 pip 和 conda 的下载超时设置。pip 可以这样：
pip config set global.timeout 600 pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple注意，如果你装 PyTorch 时用了--index-url，那一句的优先级高于全局 index-url，所以 PyTorch 还是走官方源或你指定的源，其他包走清华源。
4.2 CUDA版本不匹配与显存不足
训练阶段最刺激的报错是RuntimeError: CUDA error: no kernel image is available for execution on the device。这个报错一眼看上去很吓人，其实是 PyTorch 预编译的 CUDA 版本和显卡驱动不匹配。比如驱动只支持 CUDA 11.x，你却装了 cu121 的 PyTorch，那么 GPU 在尝试执行内核时就会报这个错。
解决方法有两种：要么降低 PyTorch 的 CUDA 版本，换成 cu118；要么更新 NVIDIA 驱动，让驱动支持更高版本的 CUDA。我一般不推荐追高版本，除非你有充分理由。
显存不足的报错则是CUDA out of memory。3D 分割任务非常吃显存，跑不起来时优先调整这几个参数：
batch_size从 2 降到 1patch_size从 192 降到 160 或 128- 使用混合精度训练（AMP），在 nnU-Net v2 配置里开启
use_amp - 关闭可能存在的前处理缓存，或者在预处理阶段减小 spacing 重采样目标分辨率
记住一句话：环境没问题，训练跑不动，先改策略，别删代码。很多人一看到显存不足就想换卡，实际上把 patch_size 缩一圈，往往就能在现有显卡上先跑通整个流程。
4.3 环境变量、依赖缺失与版本冲突
依赖缺失的报错比较直白，比如ModuleNotFoundError: No module named 'monai'。这种情况一般是你在某个环节用了系统 Python 环境，或者安装顺序跳过了某些依赖。我建议回到第 3 节，确认当前which python指向umamba环境，再重新执行pip install -e .，setup.py 会自动补拉依赖。
还有一些人对“版本冲突”比较头疼，报错形如PackageNotFoundError、Conflict这种。我分享一个比较实用的排查顺序：先conda list看看关键包的版本，先确认 PyTorch 和 torchvision 版本一致；再看 MONAI 版本，U-Mamba 训练时会用到 MONAI 的RandCropByPosNegLabel、RandRotate90等增强方法，新老版本之间 API 有差异。遇到具体报错时，优先去 U-Mamba 仓库的 issues 里搜索报错行号，通常能找到官方解法，比我在这里猜版本更准。
我把自己踩过的问题整理成一个速查表，方便你对比排查：
| 症状 | 常见原因 | 解决办法 |
|---|---|---|
| conda create 卡在 Solving environment | conda 默认解析慢 | 用 mamba 创建，或换 conda-forge/镜像源 |
| pip install 速度极慢 | 默认源延迟高 | 配置清华 pypi 镜像，加大 timeout |
| torch.cuda.is_available() 为 False | torch 装了 CPU 版，或驱动版本不够 | 卸载后按驱动重装带 CUDA 的 torch |
| CUDA no kernel image is available | torch 的 CUDA 版本高于驱动支持 | 换 cu118 或更新驱动 |
| CUDA out of memory | patch/batch 过大 | 降 batch_size、patch_size，开 AMP |
| FileNotFoundError 指向 nnUNet_raw | 环境变量未配置 | 设置三个 nnUNet 路径并 source |
| 运行训练脚本时找不到命令 | 可编辑安装失败或 PATH 异常 | 重新 pip install -e .，检查which nnUNetv2_train |
排查的主线思路其实只有两条：第一，确认你当前激活的环境是umamba，不是系统环境；第二，确认 PyTorch 的 CUDA 通路是通的。环境问题九成都能用这两条解决，剩下的一成，基本都能在仓库 issues 里翻到答案。
我个人在实际操作中最受益的一个习惯是：每完成一个关键步骤，就立刻用conda env list和python -c "import ..."做一次状态检查，不要攒到最后一起验证。一开始我也图快，一条命令装完所有依赖，结果出了问题根本不知道是哪个包引起的。后来改成“装一步、验一步”，反而整体用时更短。
U-Mamba 环境创建这件事，技术动作并不复杂，复杂的是对版本兼容性和依赖顺序的把控。只要你理解了它依赖 nnU-Net 这个事实，就掌握了环境搭建的逻辑主线：先建干净环境、再装 PyTorch、然后装 nnU-Net、最后装 U-Mamba，并且始终用当前 conda 环境下对应的 Python 解释器去执行。把这套流程跑通之后，后面才是数据准备、训练规划和结果复现的硬仗。如果这篇文章能帮你在环境关卡少熬一个通宵，那这次梳理也算值了。
