在 GPU 上加速 TF/Keras LSTM 文本生成？

Question

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%).

Any suggestions on how to speed up text generation? A quick look at cprofiler shows that 90% of the time is spent on the single line predictions = model(input_eval) , so I don't think there are a lot of gains to be had elsewhere.

Also, the Tensorflow/Keras documentation https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict recommends calling the function just as is done below...

this method is designed for performance in large scale inputs. For small amount of inputs that fit in one batch, directly using call is recommended for faster execution, eg, model(x), or model(x, training=False)

Any suggestions on how to speed up text generation? Would it be possible to better use the GPU by generating multiple lines at the same time?

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))

Answer 1

To speed up the processing, I have two suggestions,

1. As you have GPU support, you may want to set unroll=True of the GRU layer. As per the Keras GRU documentation , setting unroll=True reduces some computation by using some extra memory. As your GPU consumption is quite less, you may want to use unroll=True . Using this setting, you may notice up to 2x speed boost (depending on the circumstances). However, you should avoid using unroll if the input sequence is too long.

2. I noticed that the text-generation architecture you linked uses GRU layer before a Dense layer. The GRU is given a parameter return_sequences=True . This causes the GRU layer to pass unnecessary output values to the following Dense layers and requires more computation. Generally, return_sequences=True should be only set if the following layer of the model is also an RNN layer. Therefore, try setting the parameter return_sequences=False . This may also improve performance.

Finally, the model(x, training=False) really works. I believe by maintaining these three issues, you may notice a significant performance improvement.

Answer 2

Not sure you can speed generation. You are doing num_generate forward calls on your model for an input with batch size of 1. While during training, you can operate on the whole sequence and compute a loss over it, during predicting each new character depends on the previously generated ones and the generation function doesn't run in parallel.

If you want to see higher GPU utilization, you could call predict on a batch of inputs seeded with different starting characters - this relates to your question about 'generating multiple lines at the same time'.

You could also try using the same starting character and tinker with the hidden state to input into the model, eg seeing what a randomly sampled state for the batch produces or extracting hidden state vectors for that starting character from training examples and populating the batched hidden state with those so your models goes in different directions from this initial character.