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

时间：2019-07-17 22:57:05

标签： python tensorflow keras semantic-segmentation tensorrt

我正在尝试使用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

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

我想知道是否有解决此问题的方法？ TensorRT还喜欢其他格式并支持上采样吗？还是可以用其他受支持的操作替换它？

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

提前谢谢！

3 个答案:

答案 0 :(得分：1)

嘿，我做过类似的事情，我想说，解决此问题的最佳方法是将模型this one导出到.onnx，如果您选中{{3 }}，支持upsample： support matrix for onnx

然后，您可以使用enter image description here将onnx-model转换为tensorrt，我将其转换为在pytorch中训练并具有上采样的网络。 onnx-tensorrt的仓库更加活跃，如果您选中pr标签，则可以检查其他人在编写自定义图层并从那里派生。

答案 1 :(得分：0)

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

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

答案 2 :(得分：0)

更新09/28/2019

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

原始答案：

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

最后，我们直接在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文件夹中还有一些示例。对于我的版本，这些资源很有帮助：

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

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

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

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

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

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

有用的示例：

  • TensorRT-5.1.2.2 / samples / sampleMNISTAP

第3步：加载权重

TensorRT希望权重采用archived documentation中提到的[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，所有激活图都是相同的，直到第三个块为止（两次合并后）。之后，像素值之间会有微小的差异。绝对百分比误差为10 ^ -8，所以我们认为它没有那么糟糕。我们仍在完善C ++实现的过程中。

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

相关问题
最新问题