TensorRT的PluginFormat必须是kNCHW？

Question

我的Tensorflow模型是follow (模型的一部分)。

​

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

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

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

Answer 1

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