Detect channels first/last of tensorflow saved model?

In the tensorflow documentation for tf.saved_model.load it states:

"Keras models are trackable, so they can be saved to SavedModel. The object returned by tf.saved_model.load is not a Keras object (i.e. doesn't have .fit, .predict, etc. methods). A few attributes and functions are still available: .variables, .trainable_variables and .call."

I would suggest you try to extract the number of channels using the .variables attribute and then compare with the model architecture (I assume you have some rough knowledge what the input/output size is and how many channels there should be in the first layer)

# channel last format
input_shape = (32,32,3)
# build model in keras
model = keras.Sequential(
    [
        keras.layers.InputLayer(input_shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(2, activation="softmax"),
    ]
)
model.save('model')

Then load the model with tf

loaded_model = tf.saved_model.load('model')

and get the output shape of the first layer:

loaded_model.variables[0].shape

Output:

TensorShape([3, 3, 3, 32])

If we have knowledge about the model architecture and that the output of the first layer has 32 channels, it is now clear that the model is saved in channel last. However, if you have no knowledge about the model's structure it will probably be more tricky and this solution won't suffice.

I'm sure there's a better answer than this, but if I just wanted to hack around this quickly...

If I knew the image size already, I'd try

# Arrange.

tf.saved_model.save(tf.keras.applications.MobileNet(),'/path/to/dir') 
m=tf.saved_model.load('/path/to/dir')
image_size = (224,224) # 

# Act
try:
  m(tf.zeros((1,) + image_size + (3,))
  format='channels_last'
except ValueError:
  try:
    m(tf.zeros((1,3) + image_size)
    format='channels_first'
  except ValueError:
    raise ValueError('input shape is neither None,224,224,3 nor None,3,224,224')

If I didn't know the image size, I'd consider:

1. Cycling through common sizes like 29x29, 32x32, 224x224, 256x256 and 299x299.

2. Brute force searching over all 2 ** 20 image size combinations up to 1024x1024

3. Calling m(None) and parse the str in the ValueError with a regex. It prints out the TensorSpec of the input shape:

>>> m(tf.zeros((1,1,1,3)))
Traceback (most recent call last):
[...]
ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (3 total):
    * Tensor("inputs:0", shape=(1, 1, 1, 3), dtype=float32)
    * False
    * None
  Keyword arguments: {}

Expected these arguments to match one of the following 4 option(s):

Option 1:
  Positional arguments (3 total):
    * TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
    * True
    * None
  Keyword arguments: {}

Option 2:
  Positional arguments (3 total):
    * TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
    * False
    * None
  Keyword arguments: {}

Option 3:
  Positional arguments (3 total):
    * TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
    * True
    * None
  Keyword arguments: {}

Option 4:
  Positional arguments (3 total):
    * TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
    * False
    * None
  Keyword arguments: {}