在 Tensorflow 中释放和重用 GPU
Freeing and Reusing GPU in Tensorflow
我想在 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?
是的,构建一些@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()
我想在 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?
是的,构建一些@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()