Search code examples
pythontensorflowkerasonnx

Change pre-trained model to expect channel first in TensorFlow

I have images of shape (batch, channel, height, width) = (1, 3, 224, 224) that need to be fed to a pretrained TensorFlow model. However, by default, TensorFlow expects its pretrained model input to have shape (1, 224, 224, 3).

For example:

import tensorflow as tf
import keras2onnx as k2o
import onnx

model = tf.keras.applications.MobileNetV2()
onnx_model = k2o.convert_keras(model, model.name)
onnx.save_model(onnx_model, 'mobilenetv2.onnx')

And when doing inference on the model, I later run into the following error:

InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: input_1 for the following indices
 index: 1 Got: 3 Expected: 224
 index: 3 Got: 224 Expected: 3
 Please fix either the inputs or the model.

How do I save a pretrained TensorFlow model to expect an image with channel first? Understanding of ONNX should not be necessary but is suggested for context.

Solution

  • You can change the default data format to channels first in the Keras configuration file by going to ~/.keras/keras.json and change the line that says "image_data_format": "channels_last" to "image_data_format": "channels_first".