在 TensorFlow 中释放和重用 GPU

Question

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

我想像这样的工作流程：

1. 进行 TF 计算。
2. 释放 GPU
3. 等一下
4. 第 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 计算，返回您关心的任何结果，然后关闭该进程。 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()