python - 如何加快 Tensorflow 2 keras 模型的推理速度？

Question

所以现在有一个很大的更新，从 TensorFlow 1.X 到 2.X。

在 TF 1.X 中，我习惯了一个管道，它帮助我将我的 keras 模型推向生产。
管道:keras (h5) model --> freeze & convert to pb --> optimize pb此工作流程帮助我加快推理速度，我的最终模型可以存储在单个 (pb) 文件中，而不是文件夹(请参阅 SavedModel 格式)。

如何在 TensorFlow 2.0.0 中优化我的推理模型？

我的第一印象是我需要将我的 tf.keras 模型转换为 tflite，但是由于我的 GPU 使用 float32 操作，这种转换会让我的生活更加艰难。

谢谢。

Answer 1

一种方法是使用 Tensorflow 和 TensorRT (TF-TRT) ( https://github.com/tensorflow/tensorrt ) 来优化您的模型。但是，在 Tensorflow 2 中，模型保存在文件夹中，而不是单个 .pb 文件中。 TF-TRT 优化模型也是如此，它们存储在一个文件夹中。您可以将模型转换为 TF-TRT:

from tensorflow.python.compiler.tensorrt import trt_convert as trt
converter = tf.experimental.tensorrt.Converter(input_saved_model_dir=saved_model_dir)
converter.convert() 
converter.save("trt_optimized_model") # save it to a dir

如果您要求模型需要包含在单个文件中(并且不关心 TF-TRT 提供的优化)，您可以将 SavedModel 转换为 ONNX。并使用 ONNX 运行时进行推理。您甚至可以在这里更进一步，将 ONNX 文件转换为 TensorRT ( https://developer.nvidia.com/Tensorrt )。这将为您提供一个可以使用 TensorRT 运行的优化文件(请注意，您不能再使用 Tensorflow 运行生成的文件)。