为何使用load_model加载模型预测时会出现RAM内存持续上涨问题？

1. 这是TensorFlow 2.4~2.6版本的已知Bug：通过load_model加载SavedModel格式的模型时，eager执行模式的全局上下文会默认缓存每一次推理产生的中间张量，且不会自动回收，Transformer这类存在大量中间计算张量的模型泄漏速度会更快。
2. SavedModel自带的默认推理签名会开启无限的trace缓存：如果每次输入的形状（如序列长度、batch size）存在波动，每次新形状的输入都会触发一次计算图重跟踪，每一份跟踪生成的计算图都会常驻内存，累计到一定量级就会触发OOM。
3. 直接实例化模型再加载权重不会泄漏的原因是：本地定义的模型类生成的计算图默认开启了trace缓存大小限制，且不会把中间张量绑定到全局上下文，因此内存占用稳定。

按优先级从高到低排列：

• 升级TensorFlow版本：该内存泄漏问题已经在TensorFlow 2.7及以上版本修复，升级后直接用load_model加载模型即可，无需修改其他代码。
• 若必须使用2.5版本，优先给推理函数固定输入签名，避免重复trace：

from keras.models import load_model
import tensorflow as tf

# 加载模型时添加compile=False，不需要用到训练相关的损失、优化器配置，减少冗余缓存
model = load_model(args.model, compile=False)
# 根据你的实际输入形状、数据类型修改TensorSpec参数，None代表可变维度
predict_fn = tf.function(model, input_signature=[tf.TensorSpec(shape=(None, 128), dtype=tf.int32)])

for i in range(very_big_number):
    out, _ = predict_fn(something, training=False)

• 优化你现有的临时方案，不需要重复加载完整SavedModel，直接读取权重即可：

model = myModel()
# 直接指向SavedModel目录下的权重文件前缀，不需要提前加载完整模型
model.load_weights(f"{args.model}/variables/variables")

• 极端场景下可以每轮推理后手动清理后端缓存：在循环中每次推理结束后添加tf.keras.backend.clear_session()，该方案会损失一定推理性能，仅作为备用选项。

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