TensorFlow预测阶段内存泄漏求助：CPU环境下内存无法释放

问题背景 #

训练两个MLP模型并保存权重，通过模块加载模型后执行预测，完成标签预测后内存始终无法释放，即使执行gc.collect()、tf.keras.backend.clear_session()等操作也无效，运行环境为纯CPU。

加载模型代码 #

def creat_model_extractor(model_path, feature_count):
    
    try:
        tf.keras.backend.clear_session()
        node_list = [1024, 512, 256, 128, 64, 32]

        model = Sequential()
        model.add(Input(shape=(feature_count,)))

        for node in node_list:
            model.add(Dense(node, activation='relu'))
            model.add(Dropout(0.2))
            model.add(LayerNormalization())

        model.add(Dense(16, activation='relu'))
        model.add(LayerNormalization())
        model.add(Dense(1, activation='sigmoid'))

        model.load_weights(model_path)
        model.trainable = False
        for layer in model.layers:
            layer.trainable = False
    except Exception as error:
        logger.warning(error, exc_info=True)
        return None

    return model

预测代码 #

@tf.function
def inference(model, inputs):
    return tf.stop_gradient(model(inputs, training=False))

predictions = inference(setting.SMALL_MODEL, small_blocks_normal)
small_blocks['label'] = (predictions > 0.5).numpy().astype(int)

已尝试无效操作 #

import gc
tf.keras.backend.clear_session() 
del predictions
gc.collect() 

可行解决方案 #

1. 显式删除模型实例并强制GC #

仅删除predictions无法释放模型本身占用的内存，需一并删除模型实例并触发深层垃圾回收：

# 预测完成后执行
del predictions
del setting.SMALL_MODEL  # 直接删除模型实例
tf.keras.backend.clear_session()
gc.collect()
gc.collect(2)  # 触发深层垃圾回收，清理循环引用

2. 优化tf.function的使用 #

tf.function会缓存计算图，可能导致内存无法释放，可通过两种方式优化：

• 移除@tf.function装饰器，直接在Eager模式运行（CPU环境下性能损失可忽略）

# 去掉@tf.function装饰器
def inference(model, inputs):
    return tf.stop_gradient(model(inputs, training=False))

• 添加do_not_convert装饰器，避免不必要的图缓存

@tf.function
@tf.autograph.experimental.do_not_convert
def inference(model, inputs):
    return tf.stop_gradient(model(inputs, training=False))

3. 调整模型加载时的资源清理逻辑 #

将tf.keras.backend.clear_session()移到函数最外层，确保加载新模型前彻底清理之前的图资源：

def creat_model_extractor(model_path, feature_count):
    tf.keras.backend.clear_session()  # 提前清理所有旧图资源
    try:
        node_list = [1024, 512, 256, 128, 64, 32]

        model = Sequential()
        model.add(Input(shape=(feature_count,)))

        for node in node_list:
            model.add(Dense(node, activation='relu'))
            model.add(Dropout(0.2))
            model.add(LayerNormalization())

        model.add(Dense(16, activation='relu'))
        model.add(LayerNormalization())
        model.add(Dense(1, activation='sigmoid'))

        model.load_weights(model_path)
        model.trainable = False
        for layer in model.layers:
            layer.trainable = False
    except Exception as error:
        logger.warning(error, exc_info=True)
        return None

    return model

4. 用函数作用域隔离模型生命周期 #

将模型加载、预测、清理逻辑封装到单个函数中，利用Python作用域自动释放资源：

def run_prediction(model_path, feature_count, inputs):
    # 函数内加载模型，退出作用域后自动触发资源回收
    model = creat_model_extractor(model_path, feature_count)
    predictions = model(inputs, training=False)
    result = (predictions > 0.5).numpy().astype(int)
    # 提前清理资源
    del model
    tf.keras.backend.clear_session()
    gc.collect()
    return result

# 调用方式
small_blocks['label'] = run_prediction(setting.SMALL_MODEL_PATH, feature_count, small_blocks_normal)

5. 设置TensorFlow CPU内存按需分配 #

在程序启动时配置内存策略，避免TensorFlow预先占用大量内存：

import tensorflow as tf
# 程序开头添加
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True  # 该配置对CPU内存分配同样有效
session = tf.compat.v1.Session(config=config)
tf.compat.v1.keras.backend.set_session(session)

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