Сильно увеличивается потребление памяти при использовании ELMo от Tensorflow-Hub - PullRequest
Новая
       93

Сильно увеличивается потребление памяти при использовании ELMo от Tensorflow-Hub

0 голосов
/ 07 июня 2019

Я сейчас пытаюсь сравнить сходство миллионов документов. Для первого теста на процессоре я уменьшил их примерно до 50 символов каждый и попытался получить вложение ELMo для 10 из них одновременно, как показано ниже: 

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 предложений (и даже без сохранения векторов) программа потребляет до 12 ГБ ОЗУ. Это известная проблема (другие проблемы, которые я обнаружил, предполагают что-то похожее, но не настолько экстремальное) или я ошибся?

1 Ответ

1 голос
/ 07 июня 2019

Я полагаю, что это для TensorFlow 1.x без режима Eager (иначе использование 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

...