使用 TensorRT 部署语义分割网络 (U-Net)（不支持上采样）

Question

我正在尝试使用 TensorRT 部署经过训练的 U-Net。 该模型使用 Keras（以 Tensorflow 作为后端）进行训练。 代码和这个非常相似： https : //github.com/zhixuhao/unet/blob/master/model.py

当我将模型转换为 UFF 格式时，使用如下代码：

import uff
import os
uff_fname = os.path.join("./models/", "model_" + idx + ".uff")
uff_model = uff.from_tensorflow_frozen_model(
    frozen_file = os.path.join('./models', trt_fname), output_nodes = output_names, 
    output_filename = uff_fname
)

我会收到以下警告：

Warning: No conversion function registered for layer: ResizeNearestNeighbor yet.
Converting up_sampling2d_32_12/ResizeNearestNeighbor as custom op: ResizeNearestNeighbor
Warning: No conversion function registered for layer: DataFormatVecPermute yet.
Converting up_sampling2d_32_12/Shape-0-0-VecPermuteNCHWToNHWC-LayoutOptimizer as custom op: DataFormatVecPermute

我试图通过用上采样（双线性插值）和转置卷积替换上采样层来避免这种情况。 但是转换器会给我带来类似的错误。 我检查了https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html ，似乎还不支持所有这些操作。

我想知道这个问题是否有任何解决方法？ 有没有 TensorRT 喜欢并支持上采样的任何其他格式/框架？ 或者是否可以用其他一些受支持的操作替换它？

我还在某处看到可以添加自定义操作来替换 TensorRT 那些不受支持的操作。 虽然我不太确定工作流程会如何。 如果有人可以指出自定义图层的示例，这也将非常有帮助。

先感谢您！

Answer 1

警告是因为 TensorRT 尚不支持这些操作，正如您已经提到的。 不幸的是，没有简单的方法来解决这个问题。 您要么必须修改图形（即使在训练之后）以仅使用组合支持的操作； 或者自己编写这些操作作为自定义层。

但是，有一种更好的方法可以在 C++ 中的其他设备上运行推理。 您可以将TensorFlow 与 TensorRT 混合使用。 TensorRT 将分析它支持的操作的图形并将它们转换为 TensorRT 节点，图形的其余部分将像往常一样由 TensorFlow 处理。 更多信息在这里。 这个解决方案比自己重写操作要快得多。 唯一复杂的部分是从目标设备上的源构建 TensorFlow 并生成动态库tensorflow_cc 。 最近有许多指南和支持将 TensorFlow 移植到各种架构，例如ARM 。

Answer 2

更新 09/28/2019

Nvidia 大约两周前发布了TensorRT 6.0.1 ，并添加了一个名为“IResizeLayer”的新 API。 该层支持“最近”插值，因此可用于实现上采样。 不再需要使用自定义图层/插件！

原答案：

感谢您在这里发布的所有答案和建议！

最后，我们直接在 TensorRT C++ API 中实现了网络，并从 .h5 模型文件中加载了权重。 我们还没有时间来分析和完善解决方案，但根据我们输入的测试图像，推断似乎是有效的。

这是我们采用的工作流程：

步骤 1：对上采样层进行编码。

在我们的 U-Net 模型中，所有上采样层的缩放因子为 (2, 2) 并且它们都使用 ResizeNearestNeighbor 插值。 本质上，原始张量中 (x,y) 处的像素值将变为四个像素：(2x, 2y)、(2x+1, 2y)、(2x, 2y+1) 和 (2x+1, 2y+1) ) 在新的张量中。 这可以很容易地编码到 CUDA 内核函数中。

一旦我们获得了上采样内核，我们需要用 TensorRT API 包装它，特别是IPluginV2Ext 类。 开发者参考对需要实现的功能有一些说明。 我会说 enqueue() 是最重要的函数，因为 CUDA 内核在那里执行。

TensorRT Samples 文件夹中也有示例。 对于我的版本，这些资源很有帮助：

• Github：作为自定义层的 Leaky Relu
• TensorRT-5.1.2.2/samples/sampleUffSSD
• TensorRT-5.1.2.2/samples/sampleSSD

第 2 步：使用 TensorRT API 对网络的其余部分进行编码

网络的其余部分应该非常简单。 只需从TensorRT 网络定义中找到调用不同的“addxxxLayer”函数即可。

要记住的一件事：根据您使用的 TRT 版本，添加填充的方式可能会有所不同。 我认为最新版本（5.1.5）允许开发人员在addConvolution()添加参数，以便可以选择正确的填充模式。

我的模型是使用 Keras 训练的，默认填充模式是如果填充总数不均匀，则右侧和底部获得更多填充。 有关详细信息，请查看此堆栈溢出链接。 5.1.5 中有一种模式代表这种填充方案。

如果您使用的是旧版本（5.1.2.2），则需要在卷积层之前将填充作为单独的层添加，该层有两个参数：pre-padding 和 post-padding。

此外，TensorRT 中的所有内容都是 NCHW

有用的示例：

• TensorRT-5.1.2.2/samples/sampleMNISTAP

第 3 步：加载权重

TensorRT 需要格式为 [out_c, in_c, filter_h, filter_w] 的权重，这在存档文档中提到。 Keras 的权重格式为 [filter_h, filter_w, c_in, c_out]。

我们通过在 Python 中调用model.save_weights('weight.h5')得到了一个纯权重文件。 然后我们可以使用 h5py 将权重读入 Numpy 数组，执行转置并将转置后的权重保存为新文件。 我们还使用 h5py 找出了组和数据集名称。 此信息在使用HDF5 C++ API将权重加载到 C++ 代码时使用。

我们逐层比较了 C++ 代码和 Python 代码的输出。 对于我们的 U-Net，所有的激活图都是相同的，直到第三个块（在 2 个池化之后）。 之后，像素值之间存在微小差异。 绝对百分比误差是 10^-8，所以我们认为它并没有那么糟糕。 我们仍在完善 C++ 实现。

再次感谢我们在这篇文章中得到的所有建议和答案。 希望我们的解决方案也能有所帮助！

Answer 3

嘿，我做过类似的事情，我想说解决这个问题的最好方法是将你的模型导出到.onnx ，就像这样，如果你检查onnx的支持矩阵，则支持上采样：

然后你可以使用https://github.com/onnx/onnx-tensorrt将 onnx-model 转换为 tensorrt，我用这个来转换我在 pytorch 中训练过的网络并且有上采样。 onnx-tensorrt的 repo 更活跃一些，如果您检查 pr 选项卡，您可以检查其他人在那里编写自定义层和 fork。