使用TensorRT部署语义分段网络（UNet）（无上采样支持）

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
)

我将收到以下警告：

^{pr2}$

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

我想知道有没有解决这个问题的办法？有没有其他的格式/框架，TensorRT喜欢并支持upsampling？或者可以用其他支持的操作替换它吗？在

我还看到了一些地方可以添加定制的操作来替换那些不受支持的TensorRT操作。虽然我不太确定工作流程会怎样。如果有人能指出一个自定义层的例子，这也会很有帮助。在

提前谢谢你！在

Answer 1

更新09/28/2019

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

原始答案：

感谢所有的答案和建议张贴在这里！在

最后，在TySRRT C++ API中直接实现了网络，并从H5模型文件中加载了权重。我们还没有时间对解决方案进行分析和改进，但根据我们输入的测试图像，推断似乎正在起作用。在

以下是我们采用的工作流程：

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

在我们的U-Net模型中，所有上采样层的比例因子为（2，2），它们都使用resizenearest近邻插值。本质上，原始张量中（x，y）处的像素值将变成新张量中的四个像素：（2x，2y），（2x+1，2y），（2x，2y+1）和（2x+1，2y+1）。这可以很容易地编码成一个CUDA内核函数。在

一旦我们得到了upsampling内核，我们需要用tensorrtapi包装它，特别是IPluginV2Ext class。开发人员参考有一些关于需要实现什么功能的描述。我认为enqueue（）是最重要的函数，因为CUDA内核在那里执行。在

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

• Github: Leaky Relu as custom layer
• TensorRT-5.1.2.2/样品/样品
• TensorRT-5.1.2.2/样品/样品

第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，卷积层有两个参数：前填充和后填充。在

而且，所有的东西都是NCHW在TensorRT

有用示例：

• TensorRT-5.1.2.2/样品/样品

第3步：加载重量

TensorRT希望权重的格式为[out\uc，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++实现的过程中。在

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

Answer 2

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

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

Answer 3

嘿，我也做过类似的事情，我认为解决这个问题的最好方法是用一个很好的this one将模型导出到.onnx，如果您选中support matrix for onnx，则支持upsample：

然后你可以用https://github.com/onnx/onnx-tensorrt把onnx模型转换成tensorrt，我用这个来转换我在pytorch中训练过的网络，它有upsample。onnx-tensorrt的repo更加活跃，如果您检查pr选项卡，您可以检查其他人编写自定义层并从那里派生。在