Change pre-trained model to expect channel first in TensorFlow

Question

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.

Reshaping the input tensor would not be a useful option, I assume, right? — Filippo Grazioli, Jun 20 '20 at 07:25
Right, the additional step would complicate preprocessing and comparison between models. — Joseph Konan, Jun 20 '20 at 08:13

score 2 · Accepted Answer · edited Jun 20 '20 at 07:48

2

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".

edited Jun 20 '20 at 07:48

apaderno

28,547
16
75
90

answered Jun 20 '20 at 07:18

spacecraft1013

101
4

Change pre-trained model to expect channel first in TensorFlow

1 Answers1