YOLOv5大尺寸输入训练时CUDA显存不足问题求助

问题背景 #

• 正在重要项目中训练YOLOv5 "L6"模型，数据集包含无人机拍摄图像
• 此前用RTX 3060训练"M"模型（输入尺寸640x640），车辆、着陆点等大目标检测效果良好，但人物等小目标检测表现不佳
• 为优化小目标检测，改用1280x1280输入尺寸，购入RTX 3090 TI，在已配置好DL/ML环境的WSL 2中运行代码
• 核心问题：输入尺寸超过640x640时均出现CUDA显存不足错误：
  • 运行"M6"模型（批量大小8、输入尺寸1280x1280）：显存占用仅约12GB仍报错
  • 运行"L6"模型（批量大小16、输入尺寸1280x1280）：显存占用超过24GB后立即崩溃，触发内存分配错误

错误堆栈信息 #

File "/mnt/d/Ubuntu-WSL-Workspace/Code_Space/Code Workspace/Python Projects/AI Workspace/Teknofest-AI-in-T/2023YOLOV5/Last-YOLOV5/yolov5/train.py", line 640, in <module>
    main(opt)
  File "/mnt/d/Ubuntu-WSL-Workspace/Code_Space/Code Workspace/Python Projects/AI Workspace/Teknofest-AI-in-T/2023YOLOV5/Last-YOLOV5/yolov5/train.py", line 529, in main
    train(opt.hyp, opt, device, callbacks)
  File "/mnt/d/Ubuntu-WSL-Workspace/Code_Space/Code Workspace/Python Projects/AI Workspace/Teknofest-AI-in-T/2023YOLOV5/Last-YOLOV5/yolov5/train.py", line 352, in train
    results, maps, _ = validate.run(data_dict,
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/d/Ubuntu-WSL-Workspace/Code_Space/Code Workspace/Python Projects/AI Workspace/Teknofest-AI-in-T/2023YOLOV5/Last-YOLOV5/yolov5/val.py", line 198, in run
    for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/tqdm/std.py", line 1178, in __iter__
    for obj in iterable:
  File "/mnt/d/Ubuntu-WSL-Workspace/Code_Space/Code Workspace/Python Projects/AI Workspace/Teknofest-AI-in-T/2023YOLOV5/Last-YOLOV5/yolov5/utils/dataloaders.py", line 172, in __iter__
    yield next(self.iterator)
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 634, in __next__
    data = self._next_data()
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1346, in _next_data
    return self._process_data(data)
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1372, in _process_data
    data.reraise()
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/_utils.py", line 644, in reraise
    raise exception
RuntimeError: Caught RuntimeError in pin memory thread for device 0.
Original Traceback (most recent call last):
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 34, in do_one_step
    data = pin_memory(data, device)
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 67, in pin_memory
    return [pin_memory(sample, device) for sample in data]  # Backwards compatibility.
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 67, in <listcomp>
    return [pin_memory(sample, device) for sample in data]  # Backwards compatibility.
  File "/home/yigit-ai-dev/.pyenv/versions/3.10.9/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 55, in pin_memory
    return data.pin_memory(device)
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

解决方案 #

1. 调整批量大小与梯度累积 #

• 对于L6模型，直接将批量大小从16降至8或4，配合梯度累积参数--accumulate n模拟大批量训练效果。例如批量大小设为4，--accumulate 4，等价于批量16的梯度更新频率，但显存占用仅为原来的1/4
• M6模型批量8仍报错时，先降至4尝试，同时开启梯度累积

2. 关闭Pin Memory #

错误堆栈显示问题出在pin memory线程，WSL2环境下pin memory可能导致额外显存占用。在训练命令中添加--pin-memory 0，或修改dataloader配置，关闭pin_memory选项

3. 启用自动混合精度训练 #

添加--amp参数，自动混合精度能大幅降低显存占用，且几乎不影响模型精度，对大输入尺寸场景效果显著

4. 优化WSL2显存分配 #

WSL2默认可能限制GPU显存使用或显存回收不及时：

• 在Windows用户目录下创建.wslconfig文件，添加以下配置（根据主机内存调整）：

  [wsl2]
  memory=32GB
  localhostForwarding=true
  

• 执行wsl --shutdown重启WSL，确保显存能被充分利用且及时回收

5. 模型轻量化调整 #

若上述方法仍不足，可尝试：

• 对L6模型使用--depth 0.75或--width 0.75降低模型深度/宽度，减少显存占用
• 保留1280x1280输入的同时，平衡模型复杂度与小目标检测效果

6. 清理显存碎片 #

训练前手动清理显存：

import torch
torch.cuda.empty_cache()

或在训练命令前添加export PYTORCH_NO_CUDA_MEMORY_CACHING=1，禁用PyTorch显存缓存，减少碎片

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