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

【问题标题】：Speeding up TF/Keras LSTM text generation on GPU?在 GPU 上加速 TF/Keras LSTM 文本生成？
【发布时间】：2020-05-18 17:21:40
【问题描述】：

用于文本生成的 tensorflow 官方示例 (https://github.com/tensorflow/docs/blob/master/site/en/tutorials/text/text_generation.ipynb) 在如下定义的循环中运行。文本生成感觉很慢，根据 NVTOP 的说法，它只使用了可用 GPU 资源的一小部分 (15-20%)。

关于如何加快文本生成的任何建议？快速浏览一下 cprofiler 会发现 90% 的时间都花在了 predictions = model(input_eval) 的单行上，所以我认为其他地方不会有很多收获。

此外，Tensorflow/Keras 文档 https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict 建议调用该函数，如下所示...

此方法专为大规模输入的性能而设计。为了适合一批的少量输入，直接使用 call 建议使用更快的执行速度，例如 model(x) 或 model(x, 训练=假）

关于如何加快文本生成的任何建议？是否可以通过同时生成多行来更好地使用 GPU？

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

【问题讨论】：

您是否尝试过在此generate_text 函数中添加@tf.function？您需要修复一些问题才能使其正常工作，但这将使您摆脱急切执行的状态，转而执行更好地编译的图形执行。

标签： python performance tensorflow keras

【解决方案1】：

为了加快处理速度，我有两个建议，

由于您有 GPU 支持，您可能需要设置GRU 层的unroll=True。根据 Keras GRU documentation，设置 unroll=True 通过使用一些额外的内存来减少一些计算。由于您的 GPU 消耗非常少，您可能希望使用unroll=True。使用此设置，您可能会注意到高达 2x 的速度提升（取决于具体情况）。但是，如果输入序列太长，则应避免使用展开。
我注意到linked 的文本生成架构在Dense 层之前使用了GRU 层。 GRU 被赋予参数return_sequences=True。这会导致GRU 层将不必要的输出值传递给下面的Dense 层，并且需要更多的计算。一般return_sequences=True只有在模型的下一层也是RNN层的情况下才应该设置。因此，尝试设置参数return_sequences=False。这也可以提高性能。

最后，model(x, training=False) 确实有效。我相信通过维护这三个问题，您可能会注意到性能的显着提升。

【讨论】：

如果设置unroll=True，则不会使用GRU的CUDNN实现。 CUDNN 实现比替代 GPU 实现快得多（大约 10 倍）。
是的，你是对的。在大多数情况下，CUDNN 优于一般实现。但是，我注意到在某些（我相信大多数情况下）情况下，使用“展开”可以提高性能（我不知道为什么，但我经历过）。这就是为什么我已经添加了一个短语“（视情况而定）”。
另外，我记得之前回答过这个问题。我经历了性能提升也有质疑的人。链接：stackoverflow.com/questions/62673685/…
实际上，在我的情况下，使用 unroll=True 可以提高大约 50% 的速度。我正在运行一个 conv 网络，然后是一个 fc lstm 层。

【解决方案2】：

不确定是否可以加快生成速度。您正在对您的模型进行num_generate 前向调用，以获取批量大小为 1 的输入。在训练期间，您可以对整个序列进行操作并计算其损失，在预测每个新字符期间取决于先前生成的字符和生成函数不并行运行。

如果您想看到更高的 GPU 利用率，您可以在一批以不同起始字符为种子的输入上调用 predict - 这与您关于“同时生成多行”的问题有关。

您也可以尝试使用相同的起始字符并修改隐藏状态以输入到模型中，例如查看批次的随机采样状态会产生什么，或从训练示例中提取该起始字符的隐藏状态向量，并使用这些填充批处理隐藏状态，以便您的模型从该初始字符向不同方向前进。

【讨论】：