MLP模型预测后内存膨胀致RAM耗尽，如何解决内存占用问题？

问题描述 #

训练两个MLP模型并保存权重，开发模型加载模块后，执行预测时模型占用的RAM持续增长，最终完全占满，即使单个模块运行也会出现该问题。

模型加载模块代码 #

def creat_model_extractor(model_path, feature_count):
    """
    This function create model and set weights
    :param model_path: address of weights files
    :param feature_count: Number of nodes in input layer
    """
    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'))

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

        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

预测代码 #

SMALL_MODEL = creat_model_extractor(MODEL_PATH_SMALL, small_blocks_count)

(SMALL_MODEL.inference_step(small_blocks_normal) > 0.5).numpy().astype(int)

解决措施 #

1. 优化@tf.function的使用，避免重复追踪计算图 #

@tf.function如果没有明确输入签名，每次传入不同形状的输入会重新生成计算图，导致内存累积。给inference_step添加固定的输入签名，确保只生成一次计算图：

# 修改加载函数中的inference_step定义
@tf.function(input_signature=[tf.TensorSpec(shape=(None, feature_count), dtype=tf.float32)])
def inference_step(inputs):
    return tf.stop_gradient(model(inputs, training=False))

2. 及时清理中间张量与触发垃圾回收 #

预测生成的NumPy数组会占用CPU内存，处理完结果后要立即删除并触发垃圾回收，避免内存堆积：

# 修改预测代码，添加内存清理逻辑
result = (SMALL_MODEL.inference_step(small_blocks_normal) > 0.5).numpy().astype(int)
# 处理result的业务逻辑...

# 清理内存
del result
import gc
gc.collect()

3. 避免重复创建模型，复用已加载的模型 #

如果多次调用creat_model_extractor创建模型，旧模型的引用未被释放会导致内存泄漏。确保只初始化一次模型，后续预测复用该实例：

# 全局或单例模式初始化模型，只执行一次
SMALL_MODEL = creat_model_extractor(MODEL_PATH_SMALL, small_blocks_count)

# 后续预测直接复用SMALL_MODEL，不再重新创建
for batch in data_batches:
    result = (SMALL_MODEL.inference_step(batch) > 0.5).numpy().astype(int)
    # 处理结果
    del result
    gc.collect()

4. 改用SavedModel格式加载模型，优化内存管理 #

将训练好的模型保存为SavedModel格式（而非仅权重），加载时使用tf.keras.models.load_model，TensorFlow会自动优化推理模式的内存占用：

# 修改模型加载函数
def creat_model_extractor(model_path, feature_count):
    try:
        tf.keras.backend.clear_session()
        # 加载SavedModel格式的模型
        model = tf.keras.models.load_model(model_path)
        model.trainable = False
        for layer in model.layers:
            layer.trainable = False
        
        @tf.function(input_signature=[tf.TensorSpec(shape=(None, feature_count), dtype=tf.float32)])
        def inference_step(inputs):
            return tf.stop_gradient(model(inputs, training=False))
        
        model.inference_step = inference_step
        return model
    except Exception as error:
        logger.warning(error, exc_info=True)
        return None

注：训练完成后，用model.save("saved_model_path")保存完整模型，而非仅保存权重。

5. 清理TensorFlow会话残留 #

如果模型不再使用，显式调用tf.keras.backend.clear_session()销毁计算图与模型资源：

# 当模型不再需要时执行
del SMALL_MODEL
tf.keras.backend.clear_session()
gc.collect()

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