运行 universal-sentence-encoder-large 在数据帧上迭代时发生内存泄漏

Question

我有 14 万个句子想要嵌入。我正在使用 TF_HUB Universal Sentence Encoder 并迭代句子（我知道这不是最好的方法但是当我尝试将超过 500 个句子输入模型时它崩溃了）。我的环境是： Ubuntu 18.04 Python 3.7.4 特遣部队 1.14 内存：16GB 处理器：i-5

我的代码是：

版本 1 我在 tf.session 上下文管理器

中迭代

embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder-large/3")
    df = pandas_repository.get_dataframe_from_table('sentences')
    with tf.compat.v1.Session() as session:
        session.run(tf.global_variables_initializer())
        session.run(tf.tables_initializer())
        sentence_embedding = None
        for i, row in df.iterrows():
            sentence = row['content']
            embeddings = embed([sentence])
            sentence_embedding = session.run(embeddings)
            df.at[i, 'embedding'] = sentence_embedding
            print('processed index:', i)

版本 2 我在每次迭代中打开和关闭一个会话

embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder-large/3")
    df = pandas_repository.get_dataframe_from_table('sentences')
    for i, row in df.iterrows():
        sentence = row['content']
        embeddings = embed([sentence])
        sentence_embedding = None
        with tf.compat.v1.Session() as session:
            session.run(tf.global_variables_initializer())
            session.run(tf.tables_initializer())
            sentence_embedding = session.run(embeddings)
            df.at[i, 'embedding'] = sentence_embedding
            print('processed index:', i)

虽然 版本 2 似乎有某种 GC 并且内存被清除了一点。它仍然超过 50 个项目并爆炸。

版本 1 只是在吞噬记忆。

arnoegw

给出的正确答案

def calculate_embeddings(dataframe, table_name):
    sql_get_sentences = "SELECT * FROM semantic_similarity.sentences WHERE embedding IS NULL LIMIT 1500"
    sql_update = 'UPDATE {} SET embedding = data.embedding FROM (VALUES %s) AS data(id, embedding) WHERE {}.id = data.id'.format(table_name, table_name)
    df = pandas_repository.get_dataframe_from_sql(sql_get_sentences) 
    with hub.eval_function_for_module("https://tfhub.dev/google/universal-sentence-encoder-large/3") as embed:    
        while len(df) >= 0:
            sentence_array = df['content'].values
            sentence_embeddings = embed(sentence_array)
            df['embedding'] = sentence_embeddings.tolist()
            values = [tuple(x) for x in df[['id', 'embedding']].values]
            pandas_repository.update_db_from_df('semantic_similarity.sentences', sql_update, values)       
            df = pandas_repository.get_dataframe_from_sql(sql_get_sentences)

我是 TF 的新手，可以使用我能得到的任何帮助。

Answer 1

您的代码使用 tf.Session，因此它属于 TF1.x 编程模型，即首先构建数据流图，然后运行重复输入并从中获取输出图。

但是您的代码与该编程模型不一致。两个版本都将 hub.Module 的新应用（调用）添加到默认的 TensorFlow 图，而不是一次应用它，并且运行对不同的输入重复使用相同的图。版本 2 不断进出 tf.Sessions，这会释放一些内存但效率很低。

请参阅我对“”的回答，了解如何在 TensorFlow 基于图形的编程模型中正确执行此操作 1.x。

即将发布的TensorFlow 2.0默认采用"eager execution"的编程模型，去掉了graphs和session，避免了这种混乱。 TensorFlow Hub 将适时针对 TF2.0 进行更新。有关您的用例的预览，请参阅 https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_text_classification.ipynb

运行 universal-sentence-encoder-large 在数据帧上迭代时发生内存泄漏

Memory leak when running universal-sentence-encoder-large itterating on dataframe

tensorflow

keras-layer

tf.keras

tensorflow-hub