Как уменьшить размер потребления GPU при использовании модели Elmo?

Я выполняю задачу НЛП, используя Elmo модель. Всякий раз, когда я загружаю Elmo модели, он занимает 15 ГБ памяти моего графического процессора. Как его уменьшить?

Ниже мой код

import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
tf.disable_eager_execution()
from tensorflow.compat.v1.keras import backend as K
sess = tf.Session()
K.set_session(sess)

elmo_model = hub.Module("https://tfhub.dev/google/elmo/2", trainable=True)
sess.run(tf.global_variables_initializer())
sess.run(tf.tables_initializer())


def ElmoEmbedding(x):
return elmo_model(inputs={
                        "tokens": tf.squeeze(tf.cast(x, tf.string)),
                        "sequence_len": tf.constant(batch_size*[maxlen])
                  },
                  signature="tokens",
                  as_dict=True)["elmo"]

И тогда я прохожу ElmoEmbedding в Lambda слой как ниже

input_text = Input(shape=(maxlen,), dtype=tf.string)

embedding = Lambda(ElmoEmbedding, output_shape=(maxlen, 1024))(input_text)

x = Bidirectional(LSTM(units=512, return_sequences=True,
                   recurrent_dropout=0.2, dropout=0.2))(embedding)

.....

Что мне нужно изменить в приведенном выше коде?

0

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *