如何在使用Elmo模型的同时减少GPU的消耗？

Question

我正在使用Elmo模型执行NLP任务。当我加载Elmo模型时，它占用了我的GPU内存的15 GB。我怎样才能减少呢？

下面是我的代码

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"]

然后，我在Lambda层中传递ElmoEmbedding，如下所示

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)

.....

，我需要在上面的代码中修改什么?

Answer 1

我知道您只是在进行推断，因此不应该需要您的模型是可训练的，因此您可以在调用trainable时将参数hub.Module设置为False。