TensorRT 的 PluginFormat 必须是 kNCHW 吗？答案

【问题标题】：TensorRT's PluginFormat has to be kNCHW?TensorRT 的 PluginFormat 必须是 kNCHW 吗？
【发布时间】：2019-09-04 16:02:56
【问题描述】：

我的 Tensorflow 模型如下（模型的一部分）。

Tensorflow 模型需要是 NHWC 格式的输入图像及其处理。在将 Tensorflow 模型转换为 TensorRT 引擎时，upsample 需要实现 Plugin。

但是 TensorRT 插件需要的格式必须是 PluginFormat::kNCHW。如果设置为 PluginFormat::kNHWC，则无法编译插件。

那么如何为这样的 TensorFlow 模型创建插件呢？

【问题讨论】：

【解决方案1】：

是的 TensorRT 插件需要是 NCHW 格式。要使用 NHWC 格式的 Tensorflow 模型，处理部分，例如在 CUDA 代码上运行的部分需要设计为处理 NCHW 格式的输入数组。然后如果 Tensorflow 模型是 NHWC 格式，则在插件的输出中重新格式化回 NHWC。

【讨论】：