Search code examples
deep-learningpytorchtensorflow-lite

Which format is preferable for tflite model NCHW or NHWC?

For gpus nchw mode is preferrable due to speed benefits, but what is the preferable mode in terms of mobile perfomance for tflite model? Now converting from pytorch to tflite yields a working nchw model but is this mode optimal?

Solution

  • The TensorFlow Lite's matrix multiplication library, for convolution 2d and so on, prefers NHWC inputs .

    The TFLite converter tries to automatically transform the given NCHW weights to the corresponding NHWC weights if the given weights are constant to perform well on mobile. If the given weights are not constants, the converter will add a transpose operator after the NCHW weights to choose the NHWC's based convolution 2d algorithm for mobile.