Hub: Universal sentence encoder speed

[nebulabug]

Hi,

I am trying the sample code provided with tensorflow hub. I want to use the sentence embedding vectors at runtime. It is consistently taking about 4 seconds on CPU for one sentence. Is there something i can do to speed it up ?

import tensorflow as tf
import tensorflow_hub as hub
import time

# Import the Universal Sentence Encoder's TF Hub module
embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder/1")

sentence = "I am a sentence for which I would like to get its embedding."
messages = [sentence]

with tf.Session() as session:
  session.run([tf.global_variables_initializer(), tf.tables_initializer()])
  t1 = time.time()
  message_embeddings = session.run(embed(messages))
  print time.time() - t1

[svsgoogle]

Embedding a sentence is quite fast, but creating the graph is rather slow. You can try:

with tf.Graph().as_default():
  embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder/1")

  sentence = "I am a sentence for which I would like to get its embedding."
  messages = [sentence]
  output = embed(messages)

  with tf.Session() as session:
    session.run([tf.global_variables_initializer(), tf.tables_initializer()])
    t1 = time.time()
    message_embeddings = session.run(output)
    print(time.time() - t1)

[svsgoogle]

(In case you didn't get notified about the edit to my previous comment)

Embedding a sentence is quite fast, but creating the graph is rather slow. You can try:

with tf.Graph().as_default():
  embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder/1")

  sentence = "I am a sentence for which I would like to get its embedding."
  messages = [sentence]
  output = embed(messages)

  with tf.Session() as session:
    session.run([tf.global_variables_initializer(), tf.tables_initializer()])
    t1 = time.time()
    message_embeddings = session.run(output)
    print(time.time() - t1)

[nebulabug]

Thanks for the code. There is a difference in the timing. Now it is milli seconds. But, there is no difference in the timing overall. That is because we moved the embed(messages) out of the session. Actually the time taking part is embed(messages). Now embed(messages) is about 3.7 sec.

Do you have any idea why embed([]) takes more time ?. Looks like we are reloading some part of the graph all the time.

[svsgoogle]

That's right - embed(messages) will build a part of the TensorFlow graph. Graph building can be slow but graph execution is typically fast.

To embed many sentences, you can either enlarge the "messages" list. Another option is to use

messages = tf.placeholder(dtype=tf.string, shape=[None])

and then you can give a value for the placeholder several times (e.g. in a loop) without incurring any graph building overhead:

message_embeddings = session.run(output, feed_dict={messages: [some_sentence]})

That will build the graph once and run it several times.

[nebulabug]

Got it. You are correct if i send large batch then it is much faster. I am wondering why it takes lot of time to build the graph for each sentence. Is it because the sentence length could change and we can't use the fixed graph structure ?

I am trying to get semantically similar queries. At runtime i have only one sentence and i can't use the batch. Let me dig further and i will update my findings on this issue and close it afterwards.

[svsgoogle]

How long it takes to build the graph depends mostly on the complexity of the encoder's graph and not really that much on your sentence (unless you have truly gigantic sentences :)

For your application, it's probably best to put a placeholder on the graph and use the feed_dict approach to run the graph on the sentence when it arrives.

For finding "similar" queries/sentences, you might also want to take a look at the answer to this question.

[nebulabug]

Thanks. Using the placeholder and using it worked. I am closing the issue.

[abbasshujah]

Hi, I am using universal sentence encoder and I am calling a function to encode a sentence using sessions. I tried using the way of doing it as suggested by svsgoogle. But since my function cannot use an already existing session. I am having trouble making it faster. nebulabug can you explain how you implemented it.

[pulkit23]

@nebulabug How did you achieve it? I am not getting the placeholder thing. Please assist. Thank you.

[vbardiovskyg]

The recommended way of running inference on a module is here: https://github.com/tensorflow/hub/blob/master/docs/common_issues.md#running-inference-on-a-pre-initialized-module

In this thread, it is briefly described in the comment above: https://github.com/tensorflow/hub/issues/40#issuecomment-384966318

[pulkit23]

@vbardiovskyg
Thanks a lot. Worked perfectly.

[JoshuaGRubin]

Thanks @svsgoogle, this works beautifully! But when _isn't_ it desired behavior? Is there a gotcha here? Maybe this should be the default behavior of embed.

That's right - embed(messages) will build a part of the TensorFlow graph. Graph building can be slow but graph execution is typically fast.

To embed many sentences, you can either enlarge the "messages" list. Another option is to use

messages = tf.placeholder(dtype=tf.string, shape=[None])

and then you can give a value for the placeholder several times (e.g. in a loop) without incurring any graph building overhead:

message_embeddings = session.run(output, feed_dict={messages: [some_sentence]})

That will build the graph once and run it several times.

[rv-ltran]

@vbardiovskyg, do you know how would the code be if we have to use tensorflow 2.0? I'm asking because in tf 2.0, there is no placeholder anymore

[monk1337]

@rv-ltran In tensorflow 2.0 Tf-hub module is quite faster than tf 1.x

You can use like this :

import tensorflow_hub as hub

embedding_fn = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

def encoder_enc(sentence):
    return embedding_fn(sentence)

[MyBruso]

did anyone get this embedding working efficiently in TF2.0? I am trying to get embedding of input texts in a loop and it is still very slow, taking ~ 5 sec for each text value.

[danielcer]

@MyBruso That's much slower than expected. Could you post a sample of your code?

[MyBruso]

@danielcer sorry for the late reply. Here is my code snippet to get USE embeddings in a loop.

```

TF2.0 environment

def get_pretrained_use_model():
#model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
model = hub.load("..\..\pretrainedmodels\use\4") # I have specifically downloaded the model to disk.
return model

def get_embeddings(input):
model = get_pretrained_use_model()
return model(input)

corpus_embeddings = get_embeddings(corpus).numpy()

Here I use these functions in a loop

for index, row in some_dataframe.iterrows():
new_df = row.to_frame()
new_df = new_df.T
text_to_be_matched = new_df.iloc[0][text_feature_name]
use_embeddings_text_to_be_matched = get_embeddings([text_to_be_matched])
use_similarity_matrix = cosine_similarity(use_embeddings_text_to_be_matched ,corpus_embeddings )
matched_record_in_corpus = get_matching_record_using_similarity_from_corpus(use_similarity_matrix, text_to_be_matched)
...
# I also need to update/generate the corpus embeddings in this loop as my corpus gets updated in this loop.
...
```

This program crashes after processing 700~1000 records. (100% memory and sometime 100% CPU)