tensorflow - Jetson NX 使用 TensorRT 优化 tensorflow 模型

Question

我正在尝试加速分割模型(unet-mobilenet-512x512)。我使用 FP16 精度模式将我的 tensorflow 模型转换为 tensorRT。而且速度比我预想的要低。
在优化之前，我在使用 .pb 卡住图进行推理时有 7FPS。在 tensorRT 优化后，我有 14FPS。
这是他们网站上 Jetson NX 的基准测试结果
你可以看到，unet 256x256 的分割速度是 146 FPS。我想，在最坏的情况下，我的 unet512x512 的速度应该慢 4 倍。

这是我使用 TensorRt 优化 tensorflow 保存模型的代码:

import numpy as np
from tensorflow.python.compiler.tensorrt import trt_convert as trt
import tensorflow as tf

params = trt.DEFAULT_TRT_CONVERSION_PARAMS
params = params._replace(
    max_workspace_size_bytes=(1<<32))
params = params._replace(precision_mode="FP16")
converter = tf.experimental.tensorrt.Converter(input_saved_model_dir='./model1', conversion_params=params)
converter.convert()

def my_input_fn():
  inp1 = np.random.normal(size=(1, 512, 512, 3)).astype(np.float32)
  yield [inp1]

converter.build(input_fn=my_input_fn)  # Generate corresponding TRT engines
output_saved_model_dir = "trt_graph2"
converter.save(output_saved_model_dir)  # Generated engines will be saved.


print("------------------------freezing the graph---------------------")


from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

saved_model_loaded = tf.saved_model.load(
    output_saved_model_dir, tags=[tf.compat.v1.saved_model.SERVING])
graph_func = saved_model_loaded.signatures[
    tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
frozen_func = convert_variables_to_constants_v2(
    graph_func)
frozen_func.graph.as_graph_def()

tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                logdir="./",
                name="unet_frozen_graphTensorRt.pb",
                as_text=False)

我下载了用于 Jetson NX 基准测试 ( https://github.com/NVIDIA-AI-IOT/jetson_benchmarks ) 的存储库，并且 unet256x256 的速度确实是 ~146FPS。但是没有优化模型的管道。
我怎样才能得到类似的结果？我正在寻找使我的模型(unet-mobilenet-512x512)的速度接近 30FPS 的解决方案也许我应该以其他方式(没有 tensorflow)运行推理或更改一些转换参数？
任何建议，谢谢

Answer 1

据我所知，您链接到的存储库使用了在后台使用 TensorRT (TRT) 的命令行工具。请注意 TensorRT与“TensorFlow 中的 TensorRT”又名 TensorFlow-TensorRT (TF-TRT) 不同这就是您在代码中使用的内容。 TF-TRT 和 TRT 模型在 Jetson 设备上的运行速度都比常规 TF 模型快，但 TF-TRT 模型仍然比 TRT 模型慢( source 1 ， source 2 )。
TRT 的缺点是需要在目标设备上完成到 TRT 的转换，并且成功实现它可能非常困难，因为有各种 TensorFlow 操作 TRT does not support (在这种情况下，您需要编写一个自定义插件或向上帝祈祷互联网上有人已经这样做了。……或者仅将 TensorRT 用于您的模型的一部分并在 TensorFlow 中进行预处理/后处理)。
基本上有两种方法可以将模型从 TensorFlow 模型转换为 TensorRT“引擎”又名“计划文件”，这两种方法都使用中间格式:

TF -> UFF -> TRT

TF -> ONNX -> TRT

在这两种情况下，graphsurgeon/onnx-graphsurgeon库可用于修改 TF/ONNX 图以实现图操作的兼容性。如上所述，可以通过 TensorRT 插件添加不受支持的操作。 (这确实是这里的主要挑战:不同的图形文件格式和不同的目标 GPU 支持不同的图形操作。)
还有第三种方式，您可以使用 TF -> Caffe -> TRT，显然还有第四种方式，您可以使用 Nvidia's Transfer Learning Toolkit (TLT) (基于 TF/Keras)和 a tool called tlt-converter 但我不熟悉它。不过，后一个链接确实提到了转换 UNet 模型。
注意涉及UFF和Caffe的路径are now deprecated并且支持将在 TensorRT 9.0 中删除，所以如果你想要一些面向 future 的东西，你可能应该选择 ONNX。话虽如此，我在网上遇到的大多数在线示例代码仍然使用 UFF，而 TensorRT 9.0 还需要一段时间。
无论如何，我还没有尝试将 UNet 转换为 TensorRT，但以下存储库提供了示例代码，它们可能会让您大致了解它的工作原理:

TF -> UFF -> TRT:jkjung-avt/tensorrt_demos , NVIDIA-AI-IOT/tf_to_trt_image_classification (后者使用一点 C++)

TF -> ONNX -> TRT:tensorflow-onnx , onnx-tensorrt

Keras -> ONNX -> TRT:Nvidia blog post (这个提到将 Unet 转换为 TRT!)

请注意，即使您没有设法为您的模型完成从 ONNX 到 TRT 的转换，也可以使用 ONNX 运行时进行推理 could potentially still give you a performance gain ，尤其是当您使用 CUDA 时或 TensorRT execution provider将自动启用 provided you're on a Jetson device and running the correct ONNXRuntime build . (不过，我不确定它与 TF-TRT 或 TRT 相比如何，但它可能仍然值得一试。)
最后，为了完整起见，我还要提一下，至少我的团队一直在尝试从 TF 切换到 PyTorch 的想法，部分原因是最近 Nvidia 的支持越来越好，而且 Nvidia 员工似乎也倾向于使用 PyTorch。特别是，现在有两种不同的方法可以将模型转换为 TRT:

PyTorch -> ONNX -> TRT(由 dusty_nv 使用)

PyTorch -> TRT(通过 torch2trt 直接转换)。看来quite a few Nvidia repositories用这个。