Search code examples
tensorflowtensorrt

TensorRT's PluginFormat has to be kNCHW?

My Tensorflow model is follow (a part of the model). enter image description here

Tensorflow model needs to be in NHWC format for input image and its processing. In converting the Tensorflow model to TensorRT engine, upsample needs to implement Plugin.

But TensorRT plugin needs format has to be in PluginFormat::kNCHW. If set to PluginFormat::kNHWC, plugin can't be compiled.

So how to create plugin for such Tensorflow model?

Solution

  • Yes TensorRT plugin needs to be NCHW format. To work with Tensorflow model in NHWC format, the processing part, for example the part running on CUDA code needs to design to work on the input array in NCHW format. Then if the Tensorflow model is NHWC format, the reformat back to NHWC at plugin's output.