nnUNet训练过程中GPU和CPU利用率突降问题分析与解决

问题现象

在使用nnUNet进行医学图像分割训练时，用户遇到了一个典型问题：在最初的5-6个epoch中，GPU和CPU利用率均保持在100%左右，但随后突然降至接近0%。这一现象导致训练过程几乎停滞，epoch耗时从正常的160秒左右激增至25000秒。

环境配置

用户环境配置如下：

• 虚拟GPU环境：Ubuntu 18.04.5 LTS
• GPU：Grid T4-8Q (8GB显存)
• CPU：Intel Xeon Gold 6146 (12核)
• 内存：48GB
• 软件：PyTorch 2.2.1 + CUDA 12.1 + nnUNetv2

排查过程

1. 初步检查

用户首先排除了I/O瓶颈的可能性，通过系统监控确认wa(等待I/O的CPU时间百分比)为0。同时确认训练过程中没有其他高负载进程运行。

2. 参数调整尝试

用户尝试了多种参数调整：

• 设置nnUNet_n_proc_DA=32和OMP_NUM_THREADS=1以优化CPU使用
• 修改SimpleITK读写器中的数据类型从float32到float16
• 逐步降低nnUNet_n_proc_DA值(从32降至6再到2)

3. 系统资源监控

通过监控工具发现：

• GPU显存使用约7GB/8GB
• 系统内存使用约12GB/48GB
• 交换分区未使用
• 共享内存显示为空

4. 软件版本测试

用户尝试了：

• 官方发布的nnUNetv2版本
• 从master分支直接安装最新代码
• 重新创建conda环境并严格按顺序安装依赖

根本原因分析

经过多次测试和专家分析，问题可能源于以下几个方面：

1. 虚拟化环境限制：在虚拟GPU环境中，特别是使用GRID驱动的T4显卡时，可能存在与PyTorch多进程处理的兼容性问题。

2. 多进程管理问题：nnUNet使用Python多进程进行数据增强，在虚拟化环境中可能出现进程意外终止或通信问题。

3. 资源分配异常：虽然系统显示有足够资源，但虚拟化层的资源调度可能导致实际可用资源受限。

解决方案

临时解决方案

1. 禁用多进程数据增强：

   export nnUNet_n_proc_DA=0
   

   这会显著降低训练速度，但可以验证是否为多进程问题。

2. 修改多进程启动方法： 在Python代码中设置：

   import multiprocessing
   multiprocessing.set_start_method('spawn')
   

长期解决方案

1. 使用物理工作站：如用户最终采用的方案，在物理硬件上运行避免了虚拟化层的问题。

2. 环境配置优化：

   • 确保虚拟化环境分配足够资源
   • 检查虚拟GPU驱动兼容性
   • 考虑使用其他虚拟化方案如KVM

3. nnUNet参数调优：

   export nnUNet_n_proc_DA=11  # 对于12核CPU
   export OMP_NUM_THREADS=1
   

技术建议

1. 监控工具使用：建议同时使用nvidia-smi、htop和dstat进行全面监控，观察CPU、GPU、内存和I/O的实时变化。

2. 日志分析：启用nnUNet的详细日志模式，检查是否有进程异常退出的记录。

3. 基准测试：在投入正式训练前，先使用nnUNetTrainerBenchmark进行小规模测试验证环境稳定性。

4. 虚拟化优化：如果必须使用虚拟环境，考虑：

   • 分配更多vCPU资源
   • 启用GPU直通模式
   • 调整虚拟内存参数

总结

nnUNet在虚拟化环境中的性能问题通常与多进程管理和资源分配有关。通过本文的分析和解决方案，用户可以根据自身环境选择最适合的调试方法。对于生产环境，建议优先考虑物理硬件部署以获得最佳稳定性和性能。