Jetson NX 使用 TensorRT 优化 tensorflow 模型

Question

我正在尝试加速分割模型（unet-mobilenet-512x512）。我将我的 tensorflow 模型转换为具有 FP16 精度模式的 tensorRT。而且速度比我预想的要低。 在优化之前，我有 7FPS 的 .pb 冻结图推理。 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

您好，您可以分享您遇到的错误吗？它应该使用以下步骤：

1. 将 TensorFlow/Keras 模型转换为 .pb 文件。
2. 将 .pb 文件转换为 ONNX 格式。
3. 创建一个 TensorRT 引擎。
4. 从 TensorRT 引擎运行推理。

我不确定 Unet（我会检查），但您可能有一些 onnx 不支持的操作（请分享您的错误）。

这是 Resnet-50 的示例。

转换为 .pb：

import tensorflow as tf
import keras
from tensorflow.keras.models import Model
import keras.backend as K
K.set_learning_phase(0)

def keras_to_pb(model, output_filename, output_node_names):

   """
   This is the function to convert the Keras model to pb.

   Args:
      model: The Keras model.
      output_filename: The output .pb file name.
      output_node_names: The output nodes of the network. If None, then
      the function gets the last layer name as the output node.
   """

   # Get the names of the input and output nodes.
   in_name = model.layers[0].get_output_at(0).name.split(':')[0]

   if output_node_names is None:
       output_node_names = [model.layers[-1].get_output_at(0).name.split(':')[0]]

   sess = keras.backend.get_session()

   # The TensorFlow freeze_graph expects a comma-separated string of output node names.
   output_node_names_tf = ','.join(output_node_names)

   frozen_graph_def = tf.graph_util.convert_variables_to_constants(
       sess,
       sess.graph_def,
       output_node_names)

   sess.close()
   wkdir = ''
   tf.train.write_graph(frozen_graph_def, wkdir, output_filename, as_text=False)

   return in_name, output_node_names

# load the ResNet-50 model pretrained on imagenet
model = keras.applications.resnet.ResNet50(include_top=True, weights='imagenet', input_tensor=None, input_shape=None, pooling=None, classes=1000)

# Convert the Keras ResNet-50 model to a .pb file
in_tensor_name, out_tensor_names = keras_to_pb(model, "models/resnet50.pb", None) 

然后你需要将.pb模型转换为ONNX格式。为此，您需要安装 tf2onnx。 示例：

python -m tf2onnx.convert  --input /Path/to/resnet50.pb --inputs input_1:0 --outputs probs/Softmax:0 --output resnet50.onnx 

最后一步从 ONNX 文件创建 TensorRT 引擎：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)
def build_engine(onnx_path, shape = [1,224,224,3]):

   """
   This is the function to create the TensorRT engine
   Args:
      onnx_path : Path to onnx_file. 
      shape : Shape of the input of the ONNX file. 
  """
   with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
       builder.max_workspace_size = (256 << 20)
       with open(onnx_path, 'rb') as model:
           parser.parse(model.read())
       network.get_input(0).shape = shape
       engine = builder.build_cuda_engine(network)
       return engine

def save_engine(engine, file_name):
   buf = engine.serialize()
   with open(file_name, 'wb') as f:
       f.write(buf)
def load_engine(trt_runtime, plan_path):
   with open(engine_path, 'rb') as f:
       engine_data = f.read()
   engine = trt_runtime.deserialize_cuda_engine(engine_data)
   return engine

我建议你检查这个Pytorch TRT Unet implementation

Answer 2

据我所知，您链接到的存储库在后台使用了使用 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 中删除，所以如果你想要一些面向未来的东西，你可能应该选择 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 直接转换）。好像quitea fewNvidia repositories用这个。