TensorFlow 2.x训练多模型后如何清理GPU内存？

以下是几种能有效释放GPU内存、避免内存溢出的方法，可结合你的训练代码使用：

1. 训练后显式清理模型与会话 #

每次训练完单个模型后，主动删除模型对象、清理Keras会话并触发Python垃圾回收，能直接释放模型占用的GPU资源。

在你的循环中，每次model.fit()执行完毕后添加以下代码：

# 删除当前模型对象
del model
# 清理Keras会话，释放底层资源
tf.keras.backend.clear_session()
# 触发Python垃圾回收
import gc
gc.collect()

修改后的核心循环代码示例：

import gc  # 记得在代码开头导入gc模块

for idx in range(10):
    ep = 20
    # load X_train, y_train, X_test, y_test

    for size in units:
        # create model
        model = Model(ip, out)
        model.compile(optimizer='adam',
                 loss=tf.keras.losses.BinaryCrossentropy(),
                 metrics=[])
        model.fit(X_train, y_train, epochs=ep, batch_size=256, validation_data=(X_test, y_test))
        
        # 清理当前模型资源
        del model
        tf.keras.backend.clear_session()
        gc.collect()

2. 开启GPU内存按需分配 #

在代码开头添加GPU内存增长配置，让TensorFlow不会一次性占用全部GPU显存，而是根据模型需求动态分配，训练结束后自动释放未使用的内存：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 对每个GPU开启内存增长模式
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

3. 用上下文管理器隔离模型训练 #

将单个模型的创建、训练过程放在tf.device()上下文管理器中，进一步明确资源范围，配合前面的清理方法效果更好：

for size in units:
    with tf.device('/GPU:0'):  # 指定使用第一个GPU，根据你的设备调整
        model = Model(ip, out)
        model.compile(optimizer='adam',
                 loss=tf.keras.losses.BinaryCrossentropy(),
                 metrics=[])
        model.fit(X_train, y_train, epochs=ep, batch_size=256, validation_data=(X_test, y_test))
    
    # 离开上下文后清理资源
    del model
    tf.keras.backend.clear_session()
    gc.collect()

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