TensorFlow 2如何清除计算图并释放GPU显存？

方案1：结合变量删除+垃圾回收+会话清理（轻量场景适用） #

操作步骤：

• 单个模型训练完成后，先手动删除所有关联该模型的变量，包括模型实例、自定义层实例、训练过程中生成的中间大Tensor
• 调用Keras后端会话清理接口，再触发Python全局垃圾回收

import tensorflow as tf
import gc

# 训练完单个模型后的清理步骤
del model  # 替换为你的模型实例变量名
del custom_layer_instance  # 替换为你的自定义层实例变量名
tf.keras.backend.clear_session()
gc.collect()

注意：如果你的自定义层有外部绑定的全局变量、自定义缓存，也需要一起删除，否则会残留变量引用导致显存无法释放。

方案2：前置配置显存动态增长（可配合方案1使用） #

在所有TensorFlow操作执行前完成配置，避免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)

方案3：多进程隔离（最彻底，适配所有场景） #

如果上述方案都无法解决自定义层导致的显存残留问题，直接将每个模型的训练逻辑放到独立子进程中运行，子进程执行完毕后系统会自动回收所有显存资源，不受自定义层、计算图缓存的影响：

import multiprocessing

def train_single_model(model_config, train_data):
    # 单个模型的全部训练逻辑写在这个函数内
    # 不要在全局作用域定义模型相关变量
    model = build_custom_model(model_config)
    model.fit(train_data, epochs=10)
    return model.get_weights()

if __name__ == "__main__":
    all_model_configs = [config1, config2, config3] # 替换为你的多个模型配置列表
    train_data = load_train_data()
    for config in all_model_configs:
        p = multiprocessing.Process(target=train_single_model, args=(config, train_data))
        p.start()
        p.join() # 等待子进程训练完成、退出，显存自动释放

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