Jetson NX 使用 TensorRT 优化张量流模型

Question

Jetson NX 使用 TensorRT 优化张量流模型

Pav*_*han 6 tensorflow tensorrt nvidia-jetson

我正在尝试加速分段模型（unet-mobilenet-512x512）。我将我的tensorflow模型转换为具有FP16精度模式的tensorRT。而且速度比我预想的要低。在优化之前，我使用 .pb 冻结图进行推理时的速度为 7FPS。在tensorRT oprimization之后我有14FPS。

这是 Jetson NX 网站上的基准测试结果，
您可以看到，unet 256x256 分割速度为 146 FPS。我想，在最坏的情况下，我的unet512x512的速度应该慢4倍。

这是我使用 TensorRt 优化张量流保存模型的代码：

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)

Run Code Online (Sandbox Code Playgroud)

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

Answer 1

bal*_*alu 7

据我所知，您链接到的存储库使用在幕后使用 TensorRT (TRT) 的命令行工具。请注意，TensorRT与您在代码中使用的“TensorFlow 中的 TensorRT”又名TensorFlow -TensorRT (TF-TRT)不同。TF-TRT 和 TRT 模型在 Jetson 设备上都比常规 TF 模型运行得更快，但 TF-TRT 模型仍然比 TRT 模型慢（来源 1、来源 2）。

\n

TRT 的缺点是，到 TRT 的转换需要在目标设备上完成，并且成功实现它可能非常困难，因为有多种TRT 不支持的TensorFlow 操作（在这种情况下，您需要编写一个自定义的插件或向上帝祈祷互联网上有人已经这样做了。\xe2\x80\xa6或仅将 TensorRT 用于模型的一部分并在 TensorFlow 中进行预处理/后处理）。

\n

基本上有两种方法可以将模型从 TensorFlow 模型转换为 TensorRT“引擎”（又名“计划文件”），这两种方法都使用中间格式：

\n

TF -> UFF -> TRT
TF -> ONNX -> TRT

\n

在这两种情况下，都可以使用graphsurgeon / onnx-graphsurgeon库来修改 TF/ONNX 图以实现图操作的兼容性。如上所述，可以通过 TensorRT 插件添加不支持的操作。（这确实是这里的主要挑战：不同的图形文件格式和不同的目标 GPU 支持不同的图形操作。）

\n

还有第三种方法，您可以执行 TF -> Caffe -> TRT，显然还有第四种方法，您可以使用Nvidia 的迁移学习工具包 (TLT)（基于 TF/Keras）和一个名为tlt-converterbut I\对它不熟悉。不过，后一个链接确实提到了转换 UNet 模型。

\n

请注意，涉及 UFF 和 Caffe 的路径现已弃用，并且 TensorRT 9.0 中将删除支持，因此如果您想要面向未来的东西，您可能应该选择 ONNX。话虽这么说，我在网上遇到的大多数示例代码仍然使用 UFF，而 TensorRT 9.0 还需要一段时间。

\n

无论如何，我还没有尝试将 UNet 转换为 TensorRT，但以下存储库提供了示例代码，可能会让您了解它的原理：

\n

TF -> UFF -> TRT：jkjung-avt/tensorrt_demos、NVIDIA-AI-IOT/tf_to_trt_image_classification（后者使用了一些 C++）
\n
TF -> ONNX -> TRT：tensorflow-onnx，onnx-tensorrt
\n
Keras -> ONNX -> TRT：Nvidia 博客文章（这篇文章提到将 Unet 转换为 TRT！）
\n

\n

请注意，即使您无法成功实现模型从 ONNX 到 TRT 的转换，使用 ONNX 运行时进行推理仍可能为您带来性能提升，特别是当您使用CUDA或TensorRT时如果您在 Jetson 设备上并运行正确的 ONNXRuntime build ，则会自动启用执行提供程序。（虽然我不确定它与 TF-TRT 或 TRT 相比如何，但它可能仍然值得一试。）

\n

最后，为了完整起见，我还要提一下，至少我的团队一直在考虑从 TF 切换到 PyTorch 的想法，部分原因是 Nvidia 的支持最近变得更好，并且 Nvidia 员工似乎更倾向于 PyTorch ，也。特别是，现在有两种不同的方法可以将模型转换为 TRT：

\n

PyTorch -> ONNX -> TRT （由dusty_nv使用）
\n
PyTorch -> TRT（通过torch2trt直接转换）。似乎有不少 Nvidia 存储库都使用它。
\n

\n

我很高兴能提供帮助！我只是想告诉你我几个月前开始使用 TensorRT 时希望其他人告诉我的事情。这确实是一个混乱的生态系统，而我仍然不是专家。我刚刚用谷歌搜索了*很多*哈哈。不管怎样，请让我（以及这里的其他人）知道事情进展如何——据我所知，还没有其他人将 UNet 转换为 TensorRT，所以你的经验可能会对这里的其他人有所帮助！ (2认同)

归档时间：	4 年，8 月前
查看次数：	2650 次
最近记录：	4 年，6 月前