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

1 投票

3 回答

1226 浏览

提问于 2025-06-18 03:57

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, training=False)

有没有什么建议可以加快文本生成的速度？是否可以通过同时生成多行来更好地利用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))

3 个回答

使用不同的工具包来进行推理，比如OpenVINO，可以让你的模型运行得更快。它通过将模型转换为中间表示（IR），进行图形剪枝，并将一些操作合并到其他操作中，同时保持准确性，从而优化你的模型。然后，它在运行时使用向量化技术。

将Keras模型转换为OpenVINO其实很简单。关于如何操作的完整教程可以在这里找到。下面是一些简要步骤。

安装OpenVINO

最简单的方法是使用PIP命令。或者，你也可以使用这个工具来找到最适合你的安装方式。

pip install openvino-dev[tensorflow2]

将模型保存为SavedModel

OpenVINO无法转换HDF5格式的模型，所以你需要先将其保存为SavedModel格式。

import tensorflow as tf
from custom_layer import CustomLayer
model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})
tf.saved_model.save(model, 'model')

使用模型优化器转换SavedModel模型

模型优化器是OpenVINO开发包中的一个命令行工具。它将Tensorflow模型转换为IR，这是OpenVINO的默认格式。你还可以尝试使用FP16精度，这样可以在不显著降低准确性的情况下提高性能（更改数据类型）。在命令行中运行：

mo --saved_model_dir "model" --data_type FP32 --output_dir "model_ir"

运行推理

转换后的模型可以被运行时加载，并为特定设备编译，比如CPU或GPU（集成在你的CPU中，比如Intel HD Graphics）。如果你不确定哪个选择最好，可以使用AUTO。你关注的是延迟，所以我建议添加一个性能提示（如下所示），以使用满足你需求的设备。

# Load the network
ie = Core()
model_ir = ie.read_model(model="model_ir/model.xml")
compiled_model_ir = ie.compile_model(model=model_ir, device_name="AUTO", config={"PERFORMANCE_HINT":"LATENCY"})

# Get output layer
output_layer_ir = compiled_model_ir.output(0)

# Run inference on the input image
result = compiled_model_ir([input_image])[output_layer_ir]

免责声明：我在OpenVINO工作。OpenVINO是为Intel硬件优化的。

回答于 2025-06-18 由 Python大师

分享举报

不太确定你能否加快生成速度。你现在对模型进行的是num_generate次前向调用，而输入的批量大小是1。在训练的时候，你可以处理整个序列并计算损失，但在预测时，每生成一个新字符都依赖于之前生成的字符，而且生成的过程不能并行进行。

如果你想看到更高的GPU利用率，可以尝试对一批输入进行预测，每个输入用不同的起始字符。这和你问的“同时生成多行”的问题有关。

你也可以试着使用相同的起始字符，并调整隐藏状态输入到模型中，比如看看随机采样的状态对这一批的影响，或者从训练示例中提取该起始字符的隐藏状态向量，然后用这些向量填充批量的隐藏状态，这样模型就能从这个初始字符出发，朝不同的方向生成。

回答于 2025-06-18 由 Python大师

分享举报

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

因为你有GPU支持，可以考虑把 GRU 层的 unroll 设置为 True。根据Keras的 GRU 文档，设置 unroll=True 可以通过使用一些额外的内存来减少一些计算量。由于你的GPU使用率比较低，所以可以试试这个设置。使用这个设置后，你可能会看到速度提升到 2x（具体情况可能有所不同）。不过，如果输入序列太长，就不建议使用这个设置。
我注意到你链接的文本生成架构在 Dense 层之前使用了 GRU 层，并且 GRU 层设置了 return_sequences=True。这会导致 GRU 层把一些不必要的输出值传递给后面的 Dense 层，从而需要更多的计算。一般来说，只有当模型的下一层也是RNN层时，才应该设置 return_sequences=True。所以，试着把这个参数设置为 return_sequences=False，这也可能会提高性能。

最后， model(x, training=False) 确实有效。我相信只要注意这三个问题，你会发现性能有明显的提升。

回答于 2025-06-18 由 Python大师

分享举报

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

相关问题：

3 个回答

撰写回答