How to convert saved model from NCHW to NHWC?

Question

I trained my NCHW model on GPU and saved the best state. I now want to make the inference on CPU, which apparently only support NHWC (I get an error mentionning that). Do I have to retrain my model with NHWC, or is there a way to convert my model ?

Answer 1

I was in the same situation, seeing errors like this when trying to run model.predict on my GPU trained model on an instance with only CPU available:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Default MaxPoolingOp only supports NHWC on device type CPU

I eventually discovered that on Intel CPUs, one can successfully apply a model to data in NCHW format so long as MKL is enabled. With pip, one can install MKL enabled tensorflow with:

pip install intel-tensorflow

You can check that it is enabled (in tensorflow 2.3) with:

tf.python._pywrap_util_port.IsMklEnabled()