如何自定义YOLOv5 DataLoader训练VisDrone数据集及解决CUDA OOM问题

一、无需修改代码的快速优化策略 #

• 降低训练批次大小（batch size）：YOLOv5默认batch size可能超出显存容量，直接在训练命令中添加--batch-size 8（可根据显存情况调整为4、2甚至1），这是最直接有效的方法。
• 缩小输入图片分辨率：VisDrone数据集图片分辨率较高，将默认的640px改为416px或320px，训练命令添加--img 416，单张图片的显存占用会大幅降低。
• 开启自动混合精度训练（AMP）：添加训练参数--amp，通过半精度计算减少显存消耗，YOLOv5部分版本默认已开启，可手动指定确保生效。
• 清理GPU残留显存：在训练代码开头执行torch.cuda.empty_cache()，清除Colab会话中之前运行残留的显存占用。
• 切换轻量型模型：将YOLOv5x/l换成YOLOv5s/m，小模型的显存占用仅为大模型的1/3至1/2，完全满足课程项目的精度需求。
• 启用梯度累积：添加训练参数--accumulate 2，每累积2个batch的梯度再进行一次参数更新，等效于提升batch size但不增加显存占用。

二、针对DataLoader的优化（无需完全自定义） #

YOLOv5原生的DataLoader已经做过优化，仅需微调部分参数即可，无需完全重写：

1. 降低数据加载线程数（num_workers）：
   打开项目中的utils/dataloaders.py，找到create_dataloader函数，将num_workers的取值从min(os.cpu_count(), batch_size if batch_size > 1 else 0, 8)修改为min(os.cpu_count(), batch_size if batch_size > 1 else 0, 4)（或2），避免CPU内存过载间接影响GPU显存。
2. 关闭内存锁定（pin_memory）：
   同样在create_dataloader函数的DataLoader初始化中，将pin_memory=True改为pin_memory=False，减少CPU内存占用，缓解显存压力。
3. 简化数据增强操作：
   打开data/hyps/hyp.scratch-low.yaml，降低mixup、cutmix等增强的概率（如mixup: 0.0），或关闭部分复杂增强，减少预处理阶段的显存消耗。

三、是否需要自定义DataLoader？ #

不需要完全自定义DataLoader。上述微调参数的方式足以解决大部分显存问题，只有当你需要实现特殊数据加载逻辑时，才需要重写相关部分。优先尝试前面的快速优化策略，再根据情况调整DataLoader参数。

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