执行numba.cuda.close()后如何复用GPU？Jupyter中TensorFlow GPU内存清理方案

问题背景 #

在Jupyter Notebook中使用TensorFlow时，希望在不重启内核（保留RAM内容）的情况下清空GPU内存。尝试使用numba.cuda.close()或device.reset()后，GPU内存虽被释放，但再次执行TensorFlow代码时内核崩溃，无法复用GPU。示例代码如下：

#---- cell 1 ----
import tensorflow as tf
from numba import cuda 

#---- cell2 ----
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]])
c = tf.matmul(a, b)

#---- cell3 ----
#device = cuda.get_current_device(); device.reset()
cuda.select_device(0); cuda.close()

问题解答 #

1. 执行cuda.select_device(0); cuda.close()后能否重新使用GPU？若可以，如何操作？ #

可以重新使用GPU，但需要重新初始化TensorFlow的CUDA上下文——因为cuda.close()会销毁当前的CUDA上下文，而TensorFlow的GPU操作依赖这个上下文。具体操作步骤：

• 首先清除TensorFlow的旧会话和残留资源：

  tf.keras.backend.clear_session()
  

• 重新配置TensorFlow的GPU内存策略（确保按需分配）：

  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

• 之后即可正常运行TensorFlow代码，比如重新执行矩阵乘法操作，不会出现内核崩溃。

2. 若该方法不可行，在ipython notebook中是否有更好的不影响RAM的GPU内存清理方案？ #

推荐以下几种方案，均无需重启内核，且保留RAM中的内容：

方案1：重置TensorFlow Eager上下文释放未使用内存

TensorFlow的Eager模式下，重置上下文可以强制释放未被使用的GPU内存：

import tensorflow as tf
import gc

# 清除旧会话
tf.keras.backend.clear_session()
# 删除无用张量引用
del a, b, c  # 替换为你实际不再需要的变量
# 强制Python垃圾回收
gc.collect()
# 重置Eager上下文
from tensorflow.python.eager import context
context._context = None
context._create_context()

方案2：启用TensorFlow内存增长策略（提前配置）

在导入TensorFlow后立即开启内存增长，让TensorFlow按需分配GPU内存，不用时自动释放：

# 在导入TensorFlow后立即执行
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

后续只需配合tf.keras.backend.clear_session()和垃圾回收，即可有效释放未使用的GPU内存，无需调用numba的CUDA操作。

方案3：手动管理张量引用

对于不再需要的张量，手动删除引用并触发垃圾回收，让TensorFlow释放对应的GPU内存：

import gc

# 删除不再使用的张量/模型变量
del a, b, c
# 强制垃圾回收
gc.collect()
# 清除TensorFlow会话
tf.keras.backend.clear_session()

注意事项 #

• 避免混用numba的CUDA操作和TensorFlow：两者共享CUDA上下文，直接关闭numba的设备会破坏TensorFlow的运行环境，导致内核崩溃。

内容的提问来源于stack exchange，提问作者deltasata