numba和tensorflow一起给出了CUDA_ERROR_OUT_OF_MEMORY

Posted 2021-04-21

我在python中使用了numba和tensorflow来计算3-d循环问题。

首先，我多次使用numba来准备输入数据，一切都很好。然后我使用tensorflow进行深度学习，一切都很好。

接下来，当我第二次回到上面的程序，即使用numba时，它给了我一个错误CUDA_ERROR_OUT_OF_MEMORY。

如果我杀死线程并重新启动内核，我可以运行它，但在第二次运行时仍然失败。

似乎问题在于张量流，但我不知道。我在代码的开头和结尾添加了numba.cuda.defer_cleanup（）但它没有多大帮助。

有谁知道如何解决这个问题？

答案

您可以尝试调整TF在其初始分配中尝试采用的可见内存分数。例如：假设您有12GB的GPU内存并且想要分配~4GB：

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

要么

config = tf.ConfigProto()
config.gpu_options.allow_growth = True 
sess = tf.Session(config=config,graph=detection_graph)

在第二种方法中，您将告诉TF允许GPU增长。这些选项可以帮助您摆脱困境