python-3.x - 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 选项来控制内存大小。