生产环境TensorFlow模型.predict内存泄漏问题及解决方法咨询

核心问题分析 #

1. 内存泄漏根源：model.predict()在循环预测时会重复创建临时计算图、保留不必要的资源，长期运行会导致内存累积泄漏。直接调用模型（model(input)）复用静态计算图，能避免这类问题。
2. 输入报错原因：模型的Embedding层仅接受整数张量，而你直接传入字符串，且原代码未将预处理逻辑整合到模型中，导致输入类型不匹配。

具体修改步骤 #

1. 整合预处理逻辑到模型

将字符串标准化、字符转整数序列的步骤作为模型的一部分，确保模型可直接接收字符串输入，同时保证计算图静态化。

2. 修正模型构建/加载逻辑

修改Model类的代码如下：

import tensorflow as tf
from tensorflow.python.keras import layers, losses


class Model:

    @staticmethod
    @tf.keras.utils.register_keras_serializable()
    def data_standardization(input_data):
        input_data = tf.strings.lower(input_data)
        return tf.strings.regex_replace(input_data, r".", r"\0 ")

    def __init__(self, model_path=None, max_features=None, embedding_dim=None, max_seq_len=None):
        if model_path is not None:
            # 加载包含预处理层的完整模型
            self.model = tf.keras.models.load_model(model_path)
        elif all([max_features, embedding_dim, max_seq_len]):
            # 构建包含预处理层的完整模型
            vectorize_layer = layers.TextVectorization(
                standardize=self.data_standardization,
                split="whitespace",
                max_tokens=max_features + 1,
                output_mode="int",
                output_sequence_length=max_seq_len
            )
            # 注意：新构建模型时，需用训练数据调用adapt来初始化映射
            # 示例：vectorize_layer.adapt(train_texts)

            self.model = tf.keras.Sequential([
                vectorize_layer,
                layers.Embedding(max_features + 1, embedding_dim),
                layers.Dropout(0.2),
                layers.Dense(256, activation="relu"),
                layers.GlobalAveragePooling1D(),
                layers.Dropout(0.2),
                layers.Dense(1),
            ])

            self.model.compile(
                optimizer="adam",
                metrics=tf.metrics.BinaryAccuracy(threshold=0.0),
                loss=losses.BinaryCrossentropy(from_logits=True),
            )

    def predict(self, text):
        # 将输入转为张量，再调用模型
        input_tensor = tf.convert_to_tensor([text], dtype=tf.string)
        return self.model(input_tensor, training=False)[0][0].numpy()

3. 关键修改点说明

• 加入TextVectorization层：将字符串标准化、拆分、转整数序列的逻辑封装到模型中，让模型直接支持字符串输入。
• 输入转张量：用tf.convert_to_tensor把字符串转为TensorFlow字符串张量，匹配模型输入要求。
• 禁用训练模式：调用模型时指定training=False，确保Dropout等层使用推理模式，同时避免不必要的梯度计算。
• 使用numpy()取结果：模型返回的是张量，转为numpy数值方便业务代码使用。

额外注意事项 #

• 如果是加载已保存的模型，保存模型时必须包含TextVectorization层，否则需要在加载后手动添加预处理逻辑。
• 新构建模型时，必须对vectorize_layer调用adapt()方法，传入训练数据让它学习字符到整数的映射。

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