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

Question

我有 140,000 个句子要为其获取嵌入。我正在使用 TF_HUB 通用句子编码器并迭代句子（我知道这不是最好的方法，但是当我尝试将 500 多个句子输入模型时它会崩溃）。 我的环境是： Ubuntu 18.04 Python 3.7.4 TF 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 编程模型，即首先构建一个数据流图，然后重复运行它，输入输入并从图中获取输出。

但是您的代码与该编程模型不太一致。两个版本都不断向默认的 TensorFlow 图添加（调用）hub.Module 的新应用程序，而不是应用一次并针对各种输入重复运行同一个图。版本 2 不断进出 tf.Sessions，这释放了一些内存但效率非常低。

请参阅我对“Strongly increasing memory consumption when using ELMo from Tensorflow-Hub”的回答，了解如何在 TensorFlow 1.x 的基于图形的编程模型中正确执行此操作。

即将发布的 TensorFlow 2.0 默认采用“急切执行”的编程模型，它取消了图形和会话，并且可以避免这种混淆。 TensorFlow Hub 将适时更新 TF2.0。如需接近您的用例的预览，请参阅https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_text_classification.ipynb