How to speed up text generation in TensorFlow reference example notebook?

[zredlined]

The tensorflow official example for text generation (https://github.com/tensorflow/docs/blob/master/site/en/tutorials/text/text_generation.ipynb) runs in a loop as defined below. The text generation feels slow, and according to NVTOP only uses a fraction of the available GPU resources (15-20%).

def generate_text(model, start_string):
  # Evaluation step (generating text using the learned model)

  # Number of characters to generate
  num_generate = 1000

  # Converting our start string to numbers (vectorizing)
  input_eval = [char2idx[s] for s in start_string]
  input_eval = tf.expand_dims(input_eval, 0)

  # Empty string to store our results
  text_generated = []

  # Low temperatures results in more predictable text.
  # Higher temperatures results in more surprising text.
  # Experiment to find the best setting.
  temperature = 1.0

  # Here batch size == 1
  model.reset_states()
  for i in range(num_generate):
      predictions = model(input_eval)
      # remove the batch dimension
      predictions = tf.squeeze(predictions, 0)

      # using a categorical distribution to predict the character returned by the model
      predictions = predictions / temperature
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()

      # We pass the predicted character as the next input to the model
      # along with the previous hidden state
      input_eval = tf.expand_dims([predicted_id], 0)

      text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))

Do you have any suggestions on how I can speed this up? Or parallelize it by generating multiple examples at the same time? A quick look at cprofiler shows that 90% of the time is spent on the single line predictions = model(input_eval), so this is where we'd most likely find a speedup. Would appreciate any advice, and happy to submit a PR if I'm able to speed it up!

System information

• I am running the TensorFlow reference text generation example: https://github.com/tensorflow/docs/blob/master/site/en/tutorials/text/text_generation.ipynb
• Tested on Debian and Google Colab (with GPU support)
• TensorFlow installed from (source or binary): Binary
  -TensorFlow version: v2.1.0-rc2-17
• Python version: 3.7.5
• CUDA/cuDNN version: Cuda compilation tools, release 10.1, V10.1.243
• GPU model and memory: NVidia Tesla T4

Describe the current behavior
Text generation works fine, but feels slow. Using NVTOP it shows only 15% GPU utilization on average.

Describe the expected behavior
Hoping to speed up text generation by better leveraging the GPU

Standalone code to reproduce the issue
This issue can be replicated by running the standard TensorFlow text generation tutorial on Google Colaboratory with GPU

Other info / logs Include any logs or source code that would be helpful to

[Saduf2019]

@zredlined
Can you please share simple stand alone code to replicate the issue or if possible share a colab gist for us to a analyse the error

[zredlined]

@Saduf2019 It is not an error, the code just does not efficiently leverage the GPU by default and I'm hoping to find some advice on speeding it up. You can run the colab here:

https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/text/text_generation.ipynb

The line I'm hoping to speed up is in the generate_text() function above:
predictions = model(input_eval)

[zredlined]

@jvishnuvardhan It seems to me the challenge is getting parallelization for the GPU, while maintaining statefulness of the LSTM to predict the next character in the sentence.

Perhaps I can batch several lines to generate at once into the model.predict() while maintaining individual LSTM state per line in the batch? Or load multiple models as workers? Any suggestions or pseudocode would be much appreciated!

[zredlined]

Any suggestions here? It would be acceptable to generate multiple texts simultaneously to more effectively use the GPU. Any insights would be appreciated

[gowthamkpr]

@zredlined Try batching several lines to generate at once into the model.predict() while maintaining individual LSTM state per line in the batch and let us know if it speeds up on no.

[zredlined]

@gowthamkpr thanks! I can't figure out how to maintain LSTM state per line in the batch. The model.predict() appears to just update a single LSTM state after processing each line in the batch. Any suggestions on how to do this?

[MarkDaoust]

I'm working on this, for other reasons, but I'll try to fix this at the same time.
It may take a little while to land, but wrapping that in a tf.function, and batching the inputs should give a good speedup.

[Saduf2019]

@zredlined
Could you please check on tf 2,4,1 and let us know if you still face this issue.

[MarkDaoust]

I got the tf.function implementation working in that tutorial.

The tf.function only runs one step at a time, so it's still not ideal.

In this commit I fixed NMT-with attention to tf.function compile the whole loop, with batched-inputs. That should be even faster.

tensorflow/docs@9e18593

That commit got rolled-back because of 2.4/2.5 incompatibilities, but I'm planning to resubmit it as soon as tf 2.5 is released.