numba 和 tensorflow 一起給出了 CUDA_ERROR_OUT_OF_MEMORY

Question

我在 python 中使用了 numba 和 tensorflow 來計算 3-d 循環問題。

首先我多次使用numba來准備輸入數據，一切都很好。 然后我用tensorflow去深度學習，一切都很好。

接下來，當我第二次返回執行上述程序時，即使用numba，它給了我一個錯誤CUDA_ERROR_OUT_OF_MEMORY。

如果我殺死線程並重新啟動內核，我可以運行它，但在第二次運行時仍然失敗。

似乎問題出在 tensorflow 上，但我不知道。 我在代碼的開頭和結尾添加了 numba.cuda.defer_cleanup() 但它沒有多大幫助。

有誰知道如何解決這個問題？

Answer 1

您可以嘗試調整 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 增長。 這些選項應該可以幫助您擺脫問題