时间戳

首页 > 解决方案 > 如何将新的 Tensorflow 层添加到 TensorRT 引擎？

tensorflow - 如何将新的 Tensorflow 层添加到 TensorRT 引擎？

问题描述

我有 tensorflow 冻结模型，从中生成 TensorRT 引擎。

我无法重新训练模型，因为我没有所有这些必需的图像。

但是 Tensorflow 流程有一些后处理层，我喜欢添加到 TensorRT 引擎中。

最好的方法是什么？

我可以使用 TensorRT 层创建插件层吗？

这些 Tensorflow 层大多在 TensorRT 中可用，如下所示。

self.tensor_heatMat_up = tf.image.resize_area(self.tensor_output[:, :, :, :19], self.upsample_size,
                                                      align_corners=False, name='upsample_heatmat')
        self.tensor_pafMat_up = tf.image.resize_area(self.tensor_output[:, :, :, 19:], self.upsample_size,
                                                     align_corners=False, name='upsample_pafmat')
        if trt_bool is True:
            smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0, 19)
        else:
            smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0)
        gaussian_heatMat = smoother.get_output()

        max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat, window_shape=(3, 3), pooling_type='MAX', padding='SAME')
        self.tensor_peaks = tf.where(tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat,
tf.zeros_like(gaussian_heatMat))

TensorRT 具有 resize_area 的 scale，conv 用于 Smoother。不确定 TensorRT 中的 tf.equal。

如何将这些层添加到 TensorRT？可以使用graphsurgeon 或UFF 模型吗？

标签： tensorflowtensorrt

解决方案

以下步骤在 C++ 中为 TensorFlow 网络添加了一个自定义插件层：

  1. 实现 IPluginV2 和 IPluginCreator 类，如：使用 C++ 为 Caffe 添加自定义层。
  2. 将 TensorFlow 操作映射到插件操作。您可以为此使用GraphSurgeon。

  3. 调用设置了预处理 -p 标志的UFF 转换器。这将生成一个 UFF 文件，其中 TensorFlow 操作被 TensorRT 插件节点替换。 convert-to-uff frozen_inference_graph.pb -p config.py -t

  4. 使用 UFF 解析器使用 TensorRT 运行预处理和转换的 UFF 文件。有关详细信息，请参阅从框架导入模型时使用自定义层。使用 TensorFlow SSD 网络进行对象检测示例说明了如何使用 C++ 添加 UFF 中不支持的自定义层。有关如何预处理图形的演示，请参见示例文件夹中的 config.py。

尽管 C++ API 是实现自定义层的首选语言；由于可以轻松访问 CUDA 和 cuDNN 等库，您还可以在 Python 应用程序中使用自定义层。您可以按照使用 Python API指南添加自定义层。

推荐阅读