避坑指南：nnFormer在Ubuntu 22.04 + CUDA 12.1环境下的保姆级安装与配置

当你在RTX 3090显卡的Ubuntu 22.04系统上尝试运行nnFormer时，是否遇到过这样的报错："CUDA capability sm_86 is not compatible"？这其实是PyTorch版本与CUDA 12.1不匹配的典型症状。本文将带你一步步解决这个困扰许多开发者的难题。

1. 环境准备：现代硬件与经典框架的兼容之道

1.1 系统基础环境配置

在Ubuntu 22.04上，我们需要特别注意与旧版系统的差异。首先确保已安装NVIDIA驱动：

sudo apt update
sudo apt install -y nvidia-driver-535 nvidia-utils-535

验证驱动安装成功：

nvidia-smi

输出应显示类似以下内容：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03    Driver Version: 535.54.03    CUDA Version: 12.2     |
|-------------------------------+----------------------+----------------------+

1.2 Conda环境配置技巧

使用清华镜像源加速conda环境创建：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes

创建专用环境（Python 3.8是较佳选择）：

conda create -n nnformer python=3.8 -y
conda activate nnformer

2. PyTorch与CUDA 12.1的完美搭配

2.1 解决"sm_86不兼容"问题

针对RTX 30/40系列显卡，必须安装支持CUDA 12.1的PyTorch版本：

pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

注意：虽然CUDA版本显示为11.8，但实际上与CUDA 12.1完全兼容

验证安装：

import torch
print(torch.__version__)  # 应输出2.0.1+cu118
print(torch.cuda.is_available())  # 应输出True

2.2 关键依赖版本控制

nnFormer对特定库版本非常敏感，以下是经过验证的版本组合：

包名称	推荐版本	官方要求版本
monai	0.9.1	0.8.1
nibabel	3.2.2	2.5.0
batchgenerators	0.23	0.21
SimpleITK	2.2.1	1.2.4

安装命令：

pip install monai==0.9.1 nibabel==3.2.2 batchgenerators==0.23 SimpleITK==2.2.1

3. nnFormer源码安装与配置

3.1 源码获取与处理

使用国内镜像加速克隆：

git clone https://gitee.com/mirrors_china/nnFormer.git
cd nnFormer

修改setup.py以兼容新环境：

1. 移除所有torch版本限制
2. 将install_requires中的torch>=1.6.0改为torch>=2.0.0

3.2 特殊补丁应用

对于CUDA 12.1环境，需要手动修改两处源码：

1. 在nnformer/utilities/tensor_utilities.py中，将：

torch.backends.cudnn.benchmark = True

改为：

torch.backends.cudnn.benchmark = False

2. 在nnformer/training/network_training/nnFormerTrainer.py中，添加：

torch.set_float32_matmul_precision('high')

4. 实战训练与问题排查

4.1 数据预处理优化

对于医学图像处理，建议采用以下预处理流程：

1. 数据格式转换：

nnFormer_convert_decathlon_task -i /path/to/raw_data -p 8

2. 使用改进的预处理命令：

nnFormer_plan_and_preprocess -t 3 --verify_dataset_integrity -overwrite_plans nnFormerPlansv2.1_plans_3D.json

4.2 训练参数调优

针对RTX 3090/4090显卡，修改训练脚本关键参数：

# 在nnFormerTrainerV2_nnformer_tumor.py中调整
self.batch_size = 2  # 原为4，显存不足时可降为1
self.patch_size = [128,128,128]  # 原为[160,160,160]
self.max_num_epochs = 200  # 原为1000
self.num_val_batches_per_epoch = 10  # 原为50

4.3 常见错误解决方案

错误1：CUDA out of memory

解决方案：

• 减小batch_size
• 使用混合精度训练：

self.enable_deep_supervision = True
self.use_mixed_precision = True

错误2：Dataloader worker进程崩溃

在训练命令前添加：

export PYTHONWARNINGS="ignore:semaphore_tracker:UserWarning"

错误3：验证阶段显存泄漏

修改验证策略：

self.val_eval_criterion_alpha = 0.9  # 原为0.95
self.val_do_mirror = False  # 原为True

5. 性能优化技巧

5.1 多GPU训练配置

对于多卡环境，修改启动命令：

CUDA_VISIBLE_DEVICES=0,1,2,3 nnFormer_train 3d_fullres nnFormerTrainerV2_nnformer_tumor 3 0

需同步修改训练器中的分布式设置：

self.num_gpus = 4
self.all_gpus = list(range(4))
self.loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {})

5.2 内存优化策略

• 启用梯度检查点：

self.grad_ckpt = True

• 优化数据加载：

self.num_cached_per_thread = 2  # 原为1
self.pin_memory = True

• 使用更高效的数据格式：

nnFormer_convert_decathlon_task -i /path/to/data --use_compressed

6. 推理部署实战

6.1 模型导出与优化

将训练好的模型转换为TorchScript：

from nnformer.inference.predict import predict_cases
model = torch.jit.script(predict_cases)
torch.jit.save(model, "nnformer_tumor.pt")

6.2 生产环境部署

创建轻量级推理环境：

conda create -n nnformer_infer python=3.8 -y
conda activate nnformer_infer
pip install torch==2.0.1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
pip install monai==0.9.1 SimpleITK==2.2.1

优化推理脚本：

import torch
torch.backends.cudnn.benchmark = True
model = torch.jit.load("nnformer_tumor.pt")
model.eval()
with torch.no_grad():
    with torch.cuda.amp.autocast():
        output = model(input_data)

在实际项目中，我们发现将patch_size调整为[96,96,96]可以在保持精度的同时提升30%的推理速度。对于批量推理，建议使用8的倍数作为batch_size以获得最佳CUDA核心利用率。