在 Tensorflow 中释放和重用 GPU

Question

I would like to free and Reuse the GPU while using Tensorflow in a jupyter notebook.

I imagine a workflow like this:

1. Make a TF calculation.
2. Free the GPU
3. Wait a while
4. Step 1. again.

This is the code i use right no. Steps 1 to 3 are working step 4 is not:

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()

If i run this code in Jupyter Notebooks my kernel just dies. Is there a alte.netiv to cuda.close() and cuda.close() that frees the GPU while not breaking TF?

Answer 1

Yes, building somewhat of what @talonmies said, do not bring numba into this whatsoever. It's basically incompatible with the TensorFlow API.

Here is a solution where you completely free the GPU. Basically, you can launch the TF computations in a separate process, return any result that you care about, and then close the process. TensorFlow notably has issues regarding freeing GPU memory.

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()

Answer 2

try this.

from numba import cuda 
device = cuda.get_current_device()
device.reset()