在 Tensorflow 中释放和重用 GPU

Question

我想在 jupyter notebook 中使用 Tensorflow 时释放和重用 GPU。

我想象的工作流程是这样的：

进行TF计算。
释放 GPU
稍等
再次执行第 1 步。

这是我使用的密码。步骤 1 到 3 有效，步骤 4 无效：

import time

import tensorflow as tf
from numba import cuda 


def free_gpu():
    device = cuda.get_current_device()
    cuda.close()

def test_calc():
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])   
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

    # Run on the GPU
    c = tf.matmul(a, b)

test_calc()
free_gpu()
time.sleep(10)
test_calc()

如果我运行 Jupyter Notebooks 中的这段代码，我的内核就会死掉。 cuda.close() 和 cuda.close() 是否有替代方案可以在不破坏 TF 的情况下释放 GPU？

Answer 1

是的，构建一些@talonmies 所说的内容，不要将 numba 带入其中。它与 TensorFlow 基本不兼容 API.

这是一个完全释放 GPU 的解决方案。基本上，您可以在单独的进程中启动 TF 计算，return 您关心的任何结果，然后关闭该进程。 TensorFlow 在释放 GPU 内存方面存在明显问题。

from multiprocessing import Process, Queue
import tensorflow as tf

def test_calc(q):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

    # Run on the GPU
    c = tf.matmul(a, b)
    q.put(c.numpy())

q = Queue()
p = Process(target=test_calc, args=(q,))
p.start()
p.join()
result = q.get()

在 Tensorflow 中释放和重用 GPU

Freeing and Reusing GPU in Tensorflow

python

jupyter-notebook

cuda

tensorflow2.0

numba