PyTorch Training在缓存图像后退出



我有一个大约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 ram): 99% 11880/12000 [00:47<00:00]它/s)

94.08

所以我解决了上面的问题。出现这个问题是因为我们预先缓存了所有的图像,以便在epoch期间提高速度。这可能会提高速度,但另一方面,它也会消耗内存。当你使用Google Colab时,它为你提供12.69GB的RAM。当缓存如此庞大的数据时,所有的RAM都被消耗掉了,没有任何东西可以缓存验证集,因此,它立即关闭。有两种基本方法可以解决这个问题:

方法1:

我只是将图像大小从800减小到640由于我的训练图像不包含任何小物体,所以我实际上不需要大尺寸的图像。它减少了我的内存消耗50%

--img 640

train:缓存图像(6.6GB ram): 100% 12000/12000 [00:30<00:00,它/s)

254.08

方法2:

我在命令的末尾写了一个参数,我用它来运行这个项目:

--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

也许你应该加上" VRAM消耗";你的头衔,因为这是你训练失败的主要原因。

你的回答仍然是正确的,但是我想更详细地了解为什么这种崩溃会发生在有这种问题的人身上。

Yolov5适用于图像大小为x32。如果您的imagesize不是x32的倍数,那么Yolov5将尝试在每个epoch中拉伸图像并消耗大量不应该消耗的VRAM(至少不是为了这个)。

大的图像大小也会消耗大量的VRAM,所以即使它是x32的倍数,您的设置或配置也不足以进行此培训。

缓存命令加速你的训练,但缺点是消耗更多的VRAM。

批量大小是VRAM消耗的一个重要因素。如果你真的想用一个大的Imagesize进行训练,你应该将你的批处理大小减少x2的倍数。

我希望这能帮助到一些人。