YOLO - tensorflow 适用于 cpu 但不适用于 gpu

Question

我使用我的 GPU - Nvidia 1060 3Gb 将 YOLO 检测与训练模型一起使用，一切正常。

现在我正在尝试使用参数 --gpu 1.0 生成我自己的模型。 Tensorflow 可以看到我的 gpu，正如我在开始时可以阅读的那样： “名称：GeForce GTX 1060 主要：6 次要：1 memoryClockRate(GHz)：1.6705” "totalMemory: 3.00GiB freeMemory: 2.43GiB"

无论如何，稍后，当程序加载数据并尝试开始学习时，我收到以下错误： “未能从设备分配 832.51M（872952320 字节）：CUDA_ERROR_OUT_OF_MEMORY”

我检查了它是否尝试使用我的其他 gpu (Intel 630)，但它没有。

当我在没有“--gpu”模式选项的情况下运行训练过程时，它工作正常，但速度很慢。 （我也试过 --gpu 0.8, 0.4 等）

知道如何解决吗？

Answer 1

问题解决了。在配置文件中更改批量大小和图像大小似乎没有帮助，因为它们没有正确加载。我必须转到 defaults.py 文件并将它们更改为更低，以便我的 GPU 可以计算步骤。

Answer 2

看起来您的自定义模型占用了大量内存，而显卡无法支持它。您只需要使用 --batch 选项来控制内存大小。