使用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

我试图通过用upsampling（双线性插值）替换upsampling层并转置卷积来避免这种情况。但是转换器会抛出类似的错误。我检查了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）的指南和支持。