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

Posted 2023-03-27

【中文标题】使用 TensorRT 部署语义分割网络（U-Net）（不支持上采样）

【英文标题】：Deploy Semantic Segmentation Network (U-Net) with TensorRT (no upsampling support)

【发布时间】：2019-11-26 19:49:40

【问题描述】：

我正在尝试使用 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 不支持的操作。虽然我不太确定工作流程会如何。如果有人可以指出自定义图层的示例，那也将非常有帮助。

提前谢谢你！

【问题讨论】：

可以连接网络吗？

好像有ResizeNearestNeighbor层，好像是自定义层。

> 转置卷积 TransposedConvolution2D 在 TensorFlow 支持的操作列表中。

【参考方案1】：

警告是因为这些操作是 TensorRT 的not supported yet，正如您已经提到的。 不幸的是，没有简单的方法可以解决这个问题。您要么必须修改图形（甚至 after training）才能仅使用支持组合的操作；或者自己将这些操作写成custom layer。

但是，有一种更好的方法可以在 C++ 中的其他设备上运行推理。您可以使用TensorFlow mixed with TensorRT together。 TensorRT 将分析它支持的操作图并将它们转换为 TensorRT 节点，图的其余部分将照常由 TensorFlow 处理。更多信息here。这个解决方案比自己重写操作要快得多。唯一复杂的部分是从目标设备上的源代码和generating the dynamic librarytensorflow_cc 构建 TensorFlow。最近有许多指南和支持将 TensorFlow 移植到各种架构，例如ARM.

【讨论】：

是的，构建它非常简单，但我在部署 TensorFlow 与 TensorRT 时遇到的问题之一是直接在 GPU 阵列上运行。似乎 vanilla TensorFlow 不会直接加载 GPU 映射指针。对于我的应用程序，这会搞砸。

我认为TensorFlow支持直接GPU内存访问like here！

【参考方案2】：

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 class。开发者参考对需要实现哪些功能进行了一些描述。我想说 enqueue() 是最重要的函数，因为 CUDA 内核在那里执行。

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

Github: Leaky Relu as custom layer TensorRT-5.1.2.2/samples/sampleUffSSD TensorRT-5.1.2.2/samples/sampleSSD

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

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

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

我的模型是使用 Keras 训练的，默认的 padding 模式是，如果 padding 的总数不均匀，则右边和底部得到更多的 padding。检查此Stack Overflow link 了解详情。 5.1.5 中有一个mode 代表这种填充方案。

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

另外，TensorRT 中的所有东西都是 NCHW

有用的示例：

TensorRT-5.1.2.2/samples/sampleMNISTAP

第 3 步：加载权重

TensorRT 想要 [out_c, in_c, filter_h, filter_w] 格式的权重，archived documentation 中提到了这一点。 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++ 实现。

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

【讨论】：

Heeey 只是在这里添加评论，您可以使用具有固定权重且没有偏差的ConvTranspose2D 来模拟双线性插值，我为 onnx 做了这个，并将其与双线性插值自定义层和deconv 效果更好，也更快。

【参考方案3】：

嘿，我做过类似的事情，我想说解决这个问题的最好方法是将你的模型导出到.onnx，并使用this one，如果你检查support matrix for onnx，则支持上采样：

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

【讨论】：

ps。也许您可以将模型转换为 onnx 后发布，我可以为您解析为 trt :)

您好，感谢您的回答。我有一个简短的问题。 TensorRT 文档说 tensorRT 引擎是特定于设备的 (devtalk.nvidia.com/default/topic/1030042/jetson-tx1/…)。那么为了让它工作，我还必须在部署平台上运行这个脚本吗？谢谢！

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