tensorflow - TensorRT 的 PluginFormat 必须是 kNCHW 吗？

Question

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

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

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

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

Answer 1

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