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

我试图通过用上采样(双线性插值)和转置卷积替换上采样层来避免这种情况。但是转换器会给我带来类似的错误。我查了https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html似乎还不支持所有这些操作。

我想知道这个问题是否有任何解决方法？ TensorRT 是否有任何其他格式/框架喜欢并支持上采样？或者是否可以用其他一些支持的操作来替换它？

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

提前谢谢你!

Answer 1

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

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