启动nnUNet训练时突发CUDA内存不足错误求助

问题描述 #

此前已使用nnUNet成功训练多个模型，但在Jupyter Lab环境、Tesla V100 GPU上启动训练时，还未进入第一个epoch就突然出现CUDA内存不足错误。完整报错信息如下：

Traceback (most recent call last):
  File "/home/viberti/miniconda3/bin/nnUNet_train", line 8, in <module>
    sys.exit(main())
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/run/run_training.py", line 179, in main
    trainer.run_training()
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/training/network_training/nnUNetTrainerV2.py", line 440, in run_training
    ret = super().run_training()
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/training/network_training/nnUNetTrainer.py", line 317, in run_training
    super(nnUNetTrainer, self).run_training()
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/training/network_training/network_trainer.py", line 464, in run_training
    l = self.run_iteration(self.tr_gen, True)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/training/network_training/nnUNetTrainerV2.py", line 247, in run_iteration
    output = self.network(data)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/network_architecture/generic_UNet.py", line 391, in forward
    x = self.conv_blocks_context[d](x)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/network_architecture/generic_UNet.py", line 142, in forward
    return self.blocks(x)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/nnunet/network_architecture/generic_UNet.py", line 68, in forward
    return self.lrelu(self.instnorm(x))
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/instancenorm.py", line 72, in forward
    return self._apply_instance_norm(input)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/modules/instancenorm.py", line 32, in _apply_instance_norm
    return F.instance_norm(
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/torch/nn/functional.py", line 2466, in instance_norm
    return torch.instance_norm(
RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 15.78 GiB total capacity; 1.71 GiB already allocated; 6.50 MiB free; 1.73 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Exception in thread Thread-6:
Traceback (most recent call last):
  File "/home/viberti/miniconda3/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/home/viberti/miniconda3/lib/python3.9/threading.py", line 910, in run
    self._target(*self._args, **self._kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/batchgenerators/dataloading/multi_threaded_augmenter.py", line 92, in results_loop
    raise RuntimeError("Abort event was set. So someone died and we should end this madness. \nIMPORTANT: "
RuntimeError: Abort event was set. So someone died and we should end this madness. 
IMPORTANT: This is not the actual error message! Look further up to see what caused the error. Please also check whether your RAM was full
Exception in thread Thread-5:
Traceback (most recent call last):
  File "/home/viberti/miniconda3/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/home/viberti/miniconda3/lib/python3.9/threading.py", line 910, in run
    self._target(*self._args, **self._kwargs)
  File "/home/viberti/miniconda3/lib/python3.9/site-packages/batchgenerators/dataloading/multi_threaded_augmenter.py", line 92, in results_loop
    raise RuntimeError("Abort event was set. So someone died and we should end this madness. \nIMPORTANT: "
RuntimeError: Abort event was set. So someone died and we should end this madness. 
IMPORTANT: This is not the actual error message! Look further up to see what caused the error. Please also check whether your RAM was full

注：nnUNet会自动适配batch size，常规调小batch size方案不适用

可行解决方案 #

• 修复CUDA内存碎片化
  报错提示内存预留远大于已分配，说明存在内存碎片化问题。在训练代码开头添加环境变量配置：

  import os
  os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
  

  该参数限制内存块的最大拆分尺寸，让零散的空闲内存能被有效利用，默认值可根据实际情况调整（比如64或256）。

• 清理Jupyter环境残留内存
  Jupyter内核可能留存之前运行的张量或模型数据，占用GPU资源：

  1. 重启当前Jupyter内核（菜单栏Kernel -> Restart Kernel）
  2. 在训练代码开头手动释放GPU缓存：

     import torch
     torch.cuda.empty_cache()
     

• 降低数据加载线程数
  nnUNet默认的多线程数据增强会占用额外CPU和GPU内存，可减少线程数：
  在训练命令中添加参数：

  nnUNet_train 3d_fullres nnUNetTrainerV2 TaskXXX 0 --num_threads 2
  

  替换TaskXXX为你的任务编号，线程数可根据机器配置调整（比如1-4）。

• 排查GPU占用进程
  使用命令查看当前GPU的进程占用情况：

  nvidia-smi
  

  如果发现无关进程占用GPU内存，使用kill -9 [进程ID]命令终止进程释放资源。

• 切换到终端环境训练
  Jupyter Lab本身会占用少量GPU内存，尝试直接在终端运行训练命令，避免交互式环境的额外开销：

  nnUNet_train 3d_fullres nnUNetTrainerV2 TaskXXX 0
  

内容的提问来源于stack exchange，提问作者Vibe98