TensorFlow中如何释放GPU内存？是否有类似allow_shrink的机制？

TensorFlow目前没有类似TF_FORCE_GPU_ALLOW_SHRINK的环境变量，受CUDA驱动机制限制，GPU显存分配后默认不会主动收缩。不过可以通过手动方式强制释放闲置显存，以下是具体方案：

一、基础手动释放方法 #

训练完单个模型后，执行以下步骤彻底释放显存：

• 删除模型的Python引用：del model（替换为你的实际模型变量名）
• 清除Keras会话资源：tf.keras.backend.clear_session()
• 触发Python垃圾回收：import gc; gc.collect()

三步结合可最大程度释放模型占用的GPU显存，避免后续小模型训练时显存被无效占用。

二、集成到Keras回调函数 #

如果需要每个模型训练完成后自动执行释放操作，可自定义Keras回调类，将释放逻辑加入on_train_end方法：

import tensorflow as tf
import gc

class GPUClearCallback(tf.keras.callbacks.Callback):
    def on_train_end(self, logs=None):
        # 删除当前模型引用
        del self.model
        # 清除Keras会话，释放底层资源
        tf.keras.backend.clear_session()
        # 强制垃圾回收，清理残留内存引用
        gc.collect()

训练模型时，将该回调传入fit方法即可：

model.fit(x_train, y_train, callbacks=[GPUClearCallback()])

三、针对Keras-Tuner的适配 #

Keras-Tuner会多次构建、训练不同模型，需在每次调优迭代后执行释放操作：

• 可在tuner.search()结束后手动执行上述释放代码；
• 也可在自定义的模型构建函数末尾添加释放逻辑（注意不要影响当前模型训练流程）；
• 若使用Tuner的回调机制（部分版本支持），可将显存释放逻辑加入调优完成后的回调中。

注意事项 #

• 执行clear_session()后，之前创建的模型、层等实例将不可再使用，确保后续代码不再依赖这些实例；
• 部分情况下，CUDA驱动可能保留少量显存缓存，属于正常现象，不会影响后续模型训练；
• 确保使用TensorFlow 2.x以上版本，即时执行模式下clear_session()的释放效果更稳定。

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