训练YOLOv11时触发torch.OutOfMemoryError问题求助

问题背景 #

使用GeForce RTX 4070 Super（11.99G显存）、PyTorch 2.5.1+cu121训练YOLOv11m目标检测模型，原训练命令运行正常：

yolo task=detect mode=train device=0 epochs=2000 batch=32 data="C:\Project\Yolo\data_custom.yaml" model="C:\Project\Yolo\yolov11m.pt" imgsz=640

为匹配原图2560×1440提升训练效果，将imgsz改为1440后执行命令：

yolo task=detect mode=train device=0 epochs=2000 batch=32 data="C:\Project\Yolo\data_custom.yaml" model="C:\Project\Yolo\yolov11m.pt" imgsz=1440

出现显存不足错误：

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.48 GiB. GPU 0 has a total capacity of 11.99 GiB of which 0 bytes is free. Of the allocated memory 19.12 GiB is allocated by PyTorch, and 406.58 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management...

可行解决方法 #

• 优化显存分配策略：设置环境变量减少显存碎片。Windows命令行先执行：

  set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  

  再运行训练命令；Linux/macOS则用export替代set。
• 降低batch size：将batch=32调整为16、8这类更小的值，也可以用batch=-1让系统自动适配当前显存支持的最大batch size。
• 启用自动混合精度训练：在训练命令中添加amp=True，该参数能大幅降低显存占用，且对训练精度影响极小。修改后命令示例：

  yolo task=detect mode=train device=0 epochs=2000 batch=16 amp=True data="C:\Project\Yolo\data_custom.yaml" model="C:\Project\Yolo\yolov11m.pt" imgsz=1440
  

• 调整图像尺寸：若需尽量接近原图尺寸，可选用1344（32的倍数，符合YOLO要求）替代1440，在保证图像信息损失较小的同时降低显存压力。
• 换用轻量化模型：暂时切换为yolov11s或yolov11n模型完成训练，后续可将训练好的权重迁移至yolov11m做微调，平衡显存占用与模型性能。
• 清理GPU显存：关闭所有占用GPU的程序（如浏览器、其他AI训练进程），重启终端或IDE后再启动训练，确保显存被完全释放。
• 关闭训练时的验证：添加val=False跳过训练过程中的验证步骤，待训练完成后单独进行验证，减少显存消耗。

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