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

Question

我正在尝试使用TensorRT部署训练好的U-Net模型。该模型是使用带有TensorFlow后端的Keras进行训练的。代码与此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不支持的操作。虽然我不太确定工作流程是什么。如果有人能够指出一个自定义层的例子，那将非常有帮助。

提前感谢您！

- Yayuchen

---

你能连接网络吗？ - Mikhail

似乎有ResizeNearestNeighbor层，这些层似乎是自定义层。 - Daniel Möller

TransposedConvolution2D是TensorFlow支持的操作列表中的一项。 - mdaoust

Answer 1

更新于2019年9月28日

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 as custom layer（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训练的，默认填充模式是如果总填充数不是偶数，则右侧和底部会获得更多填充。请查看此Stack Overflow链接获取详细信息。在5.1.5中有一个mode代表这种填充方案。

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

此外，在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找出了Group和Dataset名称。在使用HDF5 C++ API将权重加载到C ++代码中时，使用了这些信息。

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

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

- Yayuchen

---

嘿，这里只是添加一个注释，你可以使用ConvTranspose2D和固定权重以及未学习的偏置来模拟双线性插值。我在ONNX中尝试了这个方法，并将其与自定义的双线性插值层进行了比较，结果发现Deconv效果更好，速度也更快。 - bpinaya

Answer 2

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

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

- Xenon

---

是的，构建它非常容易，但我在部署TensorFlow与TensorRT时遇到的一个问题是直接在GPU数组上操作。似乎普通的TensorFlow不能直接加载GPU映射指针。对于我的应用程序，这将会搞砸一些东西。 - Mikhail

我认为TensorFlow支持直接访问GPU内存，类似于这里所示:（https://github.com/tensorflow/tensorflow/blob/5912f51d580551e5cee2cfde4cb882594b4d3e60/tensorflow/core/common_runtime/direct_session_test.cc#L1946）！ - Xenon

Answer 3

嘿，我做过类似的事情，我认为解决问题的最佳方法是使用像这个这样的工具将您的模型导出到.onnx，如果您检查onnx的支持矩阵，会发现upsample得到了支持：

然后，您可以使用https://github.com/onnx/onnx-tensorrt将onnx模型转换为tensorrt，我用它来转换在pytorch中训练的具有上采样功能的网络。 onnx-tensorrt的repo更加活跃，如果您检查pr标签，您可以查看其他人编写的自定义层并从那里派生。

- bpinaya

---

也许你可以把模型转换成ONNX格式后分享给我，我可以帮你解析成TensorRT格式。 :) - bpinaya

您好，感谢您的回答。我有一个简短的问题。TensorRT文档中提到TensorRT引擎是特定于设备的（https://devtalk.nvidia.com/default/topic/1030042/jetson-tx1/loading-of-the-tensorrt-engine-in-c-api/）。因此，为了使其工作，我需要在部署平台上运行此脚本吗？谢谢！ - Yayuchen

是的，您需要在最终使用的平台上对其进行序列化。但是onnx-trt是基于jetson和xavier构建的。 - bpinaya