如何将新的 Tensorflow 层添加到 TensorRT 引擎？答案

【问题标题】：How to addin new Tensorflow layers to TensorRT engine?如何将新的 Tensorflow 层添加到 TensorRT 引擎？
【发布时间】：2019-07-22 02:30:03
【问题描述】：

我有 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，对 Smoother 有 conv。不确定 TensorRT 中的 tf.equal。

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

【问题讨论】：

标签： tensorflow tensorrt

【解决方案1】：

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

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

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

【讨论】：