YOLOv5-PyTorch训练在图像缓存阶段异常退出问题求助

嘿，我之前在Colab上训练YOLOv5的时候也碰到过几乎一模一样的情况，咱们来拆解下问题和解决方案：

问题根源 #

从日志里的12.3GB ram就能看出来，你的训练进程已经把Colab实例的内存榨到极限了。免费版Colab的内存通常在12-16GB左右，当缓存12000张800x800的图像时，最后那1%的缓存操作直接触发了内存不足（OOM），导致程序被强制终止。

具体解决方案 #

下面几个方法按优先级排序，你可以逐个尝试：

• 改用磁盘缓存模式：把图像缓存从内存转移到磁盘，这是最直接的解决方案。修改训练命令中的--cache参数为--cache disk，这样不会占用宝贵的RAM：

  !python train.py --img 800 --batch 32 --epochs 20 --data '/content/data.yaml' --cfg ./models/custom_yolov5s.yaml --weights yolov5s.pt --name yolov5s_results --cache disk
  

  这个方法几乎不会影响训练速度，只是第一次缓存会稍微慢一点，后续训练都会复用磁盘上的缓存文件。

• 减小批量大小（batch size）：如果坚持想用内存缓存，可以把--batch 32降到--batch 16甚至--batch 8，这样能释放一部分内存给图像缓存。不过缺点是训练速度会变慢，因为每次迭代处理的图像更少。

• 降低输入图像尺寸：把--img 800改成YOLOv5默认的--img 640，更小的图像尺寸会大幅降低单张图像的内存占用，缓存整个数据集的内存需求也会跟着减少。如果你的任务对图像分辨率要求不是特别高，这个方法效果很明显。

• 切换到高内存Colab实例：如果你是Colab Pro用户，可以直接切换到「High-RAM」模式；免费版用户可以尝试重启实例，偶尔能分配到内存更大的机器，但这个方法不稳定，只能作为备选。

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