TensorFlow 2.5 如何限制GPU显存占用并释放空闲显存供FAISS使用

报错核心原因：set_memory_growth仅控制TensorFlow按需分配显存，不会主动将TensorFlow已经申请但闲置的显存释放回CUDA上下文。从你的显存日志看，TensorFlow峰值占用2.8G显存，推理结束后仅280MB显存正在被使用，剩余2.5G左右显存被TensorFlow持有但未实际利用，导致FAISS申请1G显存时触发OOM。

方案1：推理结束后主动释放TensorFlow闲置显存 #

在第二个推理流程执行完成后、迁移FAISS索引到GPU之前，插入以下代码：

import gc
import tensorflow as tf

# 清空Keras会话缓存，释放无引用的模型、张量占用
tf.keras.backend.clear_session()
# 触发Python垃圾回收，清理残留的Python层对象引用
gc.collect()
# 重置GPU显存统计（可选，方便后续排查显存占用）
tf.config.experimental.reset_memory_stats('GPU:0')

执行完成后再调用tf.config.experimental.get_memory_info('GPU:0')查看，可看到当前显存占用会回落到接近初始值，闲置的显存会被释放给CUDA供FAISS使用。

方案2：提前限制TensorFlow最大显存占用（更稳定） #

per_process_gpu_memory_fraction在TF2.x的替代方案是虚拟设备显存限制配置，该配置需要在程序最开头、所有TensorFlow操作执行前添加，优先级高于set_memory_growth：

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # 限制TensorFlow最多使用4096MB（4GB）显存，可根据实际推理需求调整大小
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]
    )
  except RuntimeError as e:
    print(e)

该方案会直接限制TensorFlow的显存申请上限，剩余显存完全预留给FAISS使用，避免出现TensorFlow提前占满显存的问题，稳定性高于事后释放。

可选优化：降低FAISS GPU显存占用 #

如果上述方案仍不足以腾出足够显存，可调整FAISS GPU资源配置，减少临时内存分配：

import faiss
res = faiss.StandardGpuResources()
# 限制FAISS GPU临时内存最大为512MB，可根据索引规模调整
res.setTempMemory(512 * 1024 * 1024)

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