TensorRT 和 TensorFlow 2

TensorRT and Tensorflow 2

我正在尝试使用 TensorRT 加速 yolov3 TF2 的推理。 我在 tensorflow 2 中使用 TrtGraphConverter 函数。

我的代码基本上是这样的:

from tensorflow.python.compiler.tensorrt import trt_convert as trt

tf.keras.backend.set_learning_phase(0)
converter = trt.TrtGraphConverter(
    input_saved_model_dir="./tmp/yolosaved/",
    precision_mode="FP16",
    is_dynamic_op=True)
converter.convert()


saved_model_dir_trt = "./tmp/yolov3.trt"
converter.save(saved_model_dir_trt)

这会产生以下错误:

Traceback (most recent call last):
  File "/home/pierre/Programs/anaconda3/envs/Deep2/lib/python3.6/site-packages/tensorflow/python/framework/importer.py", line 427, in import_graph_def
    graph._c_graph, serialized, options)  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input 1 of node StatefulPartitionedCall was passed float from conv2d/kernel:0 incompatible with expected resource.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/pierre/Documents/GitHub/yolov3-tf2/tensorrt.py", line 23, in <module>
    converter.save(saved_model_dir_trt)
  File "/home/pierre/Programs/anaconda3/envs/Deep2/lib/python3.6/site-packages/tensorflow/python/compiler/tensorrt/trt_convert.py", line 822, in save
    super(TrtGraphConverter, self).save(output_saved_model_dir)
  File "/home/pierre/Programs/anaconda3/envs/Deep2/lib/python3.6/site-packages/tensorflow/python/compiler/tensorrt/trt_convert.py", line 432, in save
    importer.import_graph_def(self._converted_graph_def, name="")
  File "/home/pierre/Programs/anaconda3/envs/Deep2/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/pierre/Programs/anaconda3/envs/Deep2/lib/python3.6/site-packages/tensorflow/python/framework/importer.py", line 431, in import_graph_def
    raise ValueError(str(e))
ValueError: Input 1 of node StatefulPartitionedCall was passed float from conv2d/kernel:0 incompatible with expected resource.

这是否意味着我的某些节点无法转换?在这种情况下,为什么我的代码在 .save 步骤中出错?

可能有点遥不可及,但是您使用的是哪个 gpu?我知道 precision_mode="FP16" 仅在某些体系结构中受支持,例如 Pascal(tx2 系列)和 Turing(~2080 系列)。使用 fp16 从 TF2 移植到 trt 时我取得了不错的成绩。

当您使用 TensorRT 时,请记住您的模型架构中可能存在不受支持的层。有TensorRT support matrix供您参考。 YOLO 包含许多未实现的自定义层,例如 "yolo layer"。

因此,如果您想将 YOLO 转换为 TensorRT 优化模型,您需要从其他方式中进行选择。

  1. 尝试 TF-TRT 优化和执行兼容的子图,允许 TensorFlow 执行剩余的图。虽然您仍然可以使用 TensorFlow 广泛而灵活的功能集,但 TensorRT 将解析模型并尽可能对图形的部分应用优化。
  2. 实施您的 custom layers with Plugin API like this 示例。

我最终使用以下代码解决了这个问题。我也从 tf 2.0.-beta0 切换到 tf-nightly-gpu-2.0-preview

params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode='FP16',
    is_dynamic_op=True)
    
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model_dir,
    conversion_params=params)
converter.convert()
saved_model_dir_trt = "/tmp/model.trt"
converter.save(saved_model_dir_trt)

感谢您的帮助

有点晚了,但我现在正在使用基于 Keras 的 YoloV4 来做这件事。我已设法将其转换为 TRT,但您如何使用转换后的模型从中进行推断?我用过:

saved_model_loaded = tf.saved_model.load(input_saved_model, tags=[tag_constants.SERVING])
signature_keys = list(saved_model_loaded.signatures.keys())
print(signature_keys)

infer = saved_model_loaded.signatures['serving_default']
print(infer.structured_outputs)

根据 NVIDIA 的 colab,但它不适用于 infer(X)。