我有一个包含大约 12k 个训练图像和 500 个验证图像的数据集。我正在使用 YOLOv5-PyTorch 来训练我的模型。当我开始训练时，当它降到 缓存图像 阶段时，它突然退出。
我用来运行它的代码如下：
!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
我正在使用 Google Colab 来训练我的模型。
这是关机前执行的命令：
train：缓存图像（12.3GB 内存）：99% 11880/12000 [00:47
【问题讨论】：
标签： pytorch google-colaboratory training-data yolov5
所以我解决了上述问题。出现问题是因为我们正在缓存所有图像，以提高 epoch 期间的速度。现在这可能会提高速度，但另一方面，它也会消耗内存。当您使用 Google Colab 时，它会为您提供 12.69GB 的 RAM。当缓存如此庞大的数据时，所有的 RAM 都被消耗了，没有任何东西可以缓存验证集，因此它立即关闭。 解决这个问题有两种基本方法：
方法一：
我只是将图像尺寸从 800 减小到 640，因为我的训练图像不包含任何小物体，所以我实际上不需要大尺寸的图像。它使我的 RAM 消耗减少了 50%
--img 640
train：缓存图像（6.6GB 内存）：100% 12000/12000 [00:30
方法二：
我在我用来运行这个项目的命令末尾写了一个参数：
--cache
此命令在第一个 epoch 中缓存整个数据集，因此可以立即再次重用它，而不是再次处理它。如果您愿意在训练速度上妥协，那么这种方法对您有用。只需简单地删除这条线，你就可以开始了。您要运行的新命令将是：
!python train.py --img 800 --batch 32 --epochs 20 --data '/content/data.yaml' --cfg ./models/custom_yolov5s.yaml --weights yolov5s.pt --name yolov5s_results
【讨论】：