如何在Jupyter中使用TensorFlow时安全释放并重复使用GPU资源

崩溃原因 #

混用numba的cuda.close()接口和TensorFlow会直接销毁当前进程的CUDA上下文，而TensorFlow初始化时会绑定该上下文，销毁后再次调用TF算子时无法重建有效上下文，直接触发内核崩溃，该接口本身就不适合用于TensorFlow的资源释放场景。

方案1：TensorFlow原生资源清理（最稳妥，无额外依赖） #

优先用TensorFlow自带的接口释放显存，搭配显存按需分配配置，完全不会破坏运行环境：

import time
import gc
import tensorflow as tf

# 必须在所有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)

def free_gpu():
    # 清理TF后端会话、缓存的算子和临时变量
    tf.keras.backend.clear_session()
    # 触发Python垃圾回收，释放已无引用的张量占用的显存
    gc.collect()

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]])
    c = tf.matmul(a, b)
    # 主动删除不再使用的变量，加速资源回收
    del a, b, c

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

方案2：子进程隔离计算（完全释放所有GPU资源） #

如果需要将GPU显存完全释放给其他进程使用，可以把TF计算放到独立子进程中执行，子进程退出后系统会自动回收所有占用的GPU资源，完全不会影响Jupyter主内核：

import time
import multiprocessing

def test_calc_wrapper():
    # 把TensorFlow导入放到子进程内部，避免主进程持有CUDA上下文引用
    import tensorflow as tf
    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)
    # 如有需要可以将计算结果返回主进程
    return c.numpy()

# 第一次执行计算
p1 = multiprocessing.Process(target=test_calc_wrapper)
p1.start()
p1.join()

time.sleep(10)

# 第二次执行计算
p2 = multiprocessing.Process(target=test_calc_wrapper)
p2.start()
p2.join()

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