如何向 TensorRT 引擎添加新的 Tensorflow 层？

Question

我有生成 TensorRT 引擎的 tensorflow 冻结模型。

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

但是 Tensorflow 进程有一些 post 处理层，我想添加到 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 的比例，平滑器的转换。在 TensorRT 中不确定 tf.equal。

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

Answer 1

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

实现 IPluginV2 和 IPluginCreator 类，如下所示：Adding A Custom Layer Using C++ For Caffe.
将 TensorFlow 操作映射到插件操作。您可以为此使用 GraphSurgeon。
调用设置了预处理 -p 标志的 UFF converter。这将生成一个 UFF 文件，其中 TensorFlow 操作替换为 TensorRT 插件节点。convert-to-uff frozen_inference_graph.pb -p config.py -t
运行 pre-processed 并使用 UFF 解析器使用 TensorRT 转换 UFF 文件。有关详细信息，请参阅 Using Custom Layers When Importing A Model From A Framework. The Object Detection With A TensorFlow SSD Network 示例说明如何使用 C++ 添加 UFF 中不支持的自定义层。有关如何 pre-process 图形的演示，请参阅示例文件夹中的 config.py。

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

如何向 TensorRT 引擎添加新的 Tensorflow 层？

How to addin new Tensorflow layers to TensorRT engine?

tensorflow

tensorrt