Я выполняю задачу НЛП, используя 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)
.....
Что мне нужно изменить в приведенном выше коде?