使用来自 Tensorflow-Hub 的 ELMo 时显着增加内存消耗

Question

我目前正在尝试比较数百万文档的相似性。对于 CPU 的第一次测试，我将它们每个减少到大约 50 个字符，并尝试像这样一次为其中的 10 个字符获取 ELMo 嵌入：

ELMO = "https://tfhub.dev/google/elmo/2"
for row in file:
    split = row.split(";", 1)
    if len(split) > 1:
        text = split[1].replace("\n", "")
            texts.append(text[:50])
    if i == 300:
        break
    if i % 10 == 0:
        elmo = hub.Module(ELMO, trainable=False)
                 executable = elmo(
                 texts,
                 signature="default",
                 as_dict=True)["elmo"]

    vectors = execute(executable)
    texts = []
    i += 1

然而，即使是这个小例子，在大约 300 个句子（甚至没有保存向量）之后，程序也会消耗高达 12GB 的 RAM。这是一个已知问题（我发现的其他问题提出了类似的问题，但不是那么极端）还是我弄错了？

Answer 1

我想这是针对没有 Eager 模式的 TensorFlow 1.x（否则使用 hub.Module 可能会遇到更大的问题）。

在该编程模型中，您需要首先在 TensorFlow 图形中表达您的计算，然后对每批数据重复执行该图形。

使用 hub.Module() 构建模块并将其应用于将输入张量映射到输出张量都是图形构建的两个部分，应该只发生一次。
输入数据的循环应该只调用 session.run() 来提供输入并从固定图形中获取输出数据。

幸运的是，已经有一个实用函数可以为您完成这一切：

import numpy as np
import tensorflow_hub as hub

# For demo use only. Extend to your actual I/O needs as you see fit.
inputs = (x for x in ["hello world", "quick brown fox"])

with hub.eval_function_for_module("https://tfhub.dev/google/elmo/2") as f:
  for pystr in inputs:
    batch_in = np.array([pystr])
    batch_out = f(batch_in)
    print(pystr, "--->", batch_out[0])

就原始 TensorFlow 而言，这对您的作用大致如下：

module = Module(ELMO_OR_WHATEVER)
tensor_in = tf.placeholder(tf.string, shape=[None])  # As befits `module`.
tensor_out = module(tensor_in)

# This kind of session handles init ops for you.
with tf.train.SingularMonitoredSession() as sess:
  for pystr in inputs:
    batch_in = np.array([pystr])
    batch_out = sess.run(tensor_out, feed_dict={tensor_in: batch_in}
    print(pystr, "--->", batch_out[0])

如果您的需求对于 with hub.eval_function_for_module ... 来说过于复杂，您可以构建这个更明确的示例。

注意 hub.Module 是如何在循环中既不构造也不调用的。

PS：厌倦了担心构建图表与运行会话？然后 TF2 和急切的执行是给你的。查看 https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_text_classification.ipynb

使用来自 Tensorflow-Hub 的 ELMo 时显着增加内存消耗

Strongly increasing memory consumption when using ELMo from Tensorflow-Hub

python

tensorflow

tensorflow-hub

elmo