Tensorflow 对象检测 API GPU 内存问题

Tensorflow Object Detection API GPU memory issue

我目前正在尝试训练一个基于模型检测动物园的模型以进行对象检测。 运行 CPU 上的设置按预期工作,但在我的 GPU 上尝试相同会导致以下错误。

2021-03-10 11:46:54.286051: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library cublas64_11.dll
2021-03-10 11:46:54.751423: E tensorflow/stream_executor/cuda/cuda_blas.cc:226] failed to create cublas handle: CUBLAS_STATUS_ALLOC_FAILED
2021-03-10 11:46:54.764147: E tensorflow/stream_executor/cuda/cuda_blas.cc:226] failed to create cublas handle: CUBLAS_STATUS_ALLOC_FAILED
2021-03-10 11:46:54.764233: W tensorflow/stream_executor/stream.cc:1455] attempting to perform BLAS operation using StreamExecutor without BLAS support

在任务管理器中监控 GPU 信息,似乎 tensorflow(据我所知是预期的)试图分配整个内存。在达到特定峰值(8Gb 中大约 7.3Gb)后不久,TF 崩溃并出现上述代码段中的错误。

互联网/Whosebug 中针对此特定错误的解决方案提到可以解决此问题,从而允许动态内存增长。这样做似乎可行,并且 TF 设法创建了至少一个新检查点,但最终因类似类别的错误而崩溃。在那种情况下,CUDA_OUT_OF_MEMORY 错误。

系统信息:

训练设置:

最初我想使用预训练的 EfficientDet D6 模型,但也尝试了其他几个模型,如 EfficientDet D4、CenterNet HourGlass 512x512 和 SSD MobileNet V2 FPNLite。所有这些模型都以不同的批量大小开始,但即使批量大小为 1,问题仍然存在。训练图像也不大(平均 600 x 800)。目前总共有30张图片,每个class 15张用于训练(我知道训练数据集应该更大,但只是为了测试设置)。

现在我的问题是,是否有人有根据的猜测或找到此错误原因的其他方法,因为我无法想象我的 2060 至少不能训练批量大小为 1 的 SSD,而是小图像。会不会是硬件故障?如果是这样,有没有办法检查?

我已经完全重新安装了所有涉及的组件。这次我可能做了一些不同的事情,但我不能说是什么。至少我现在可以使用 GPU 进行训练了。