Detect channels first/last of tensorflow saved model?

Question

Is there any way to detect the channels first or last format for TF saved model loaded as model=tf.saved_model.load(path)?

In Keras and can go over model.layers and check for a layer l l.data_format == 'channels_last'

Is there something like this for TF saved model? I can't find any suitable documentation of TF model details - everything goes back to Keras.

Answer 1

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.

Answer 2

This can be done by looking at the data_format attributes of the node-defs.

First obtain the graph-def, e.g.:

model_dir = ...
sm = tf.saved_model.load(model_dir)
func = sm.signatures['serving_default'] # or another signature depending on the graph you wish to do this for
graph_def = func.graph.as_graph_def() # you could also stick to the graph here, no need per-se to get the graph-def, see below

Within the graph-def, the attribute will be present for the nodes for which it makes sense to have it (e.g. Conv2D):

for node in graph_def.node:
    if 'data_format' in node.attr:
        layout = node.attr['data_format'].s.decode('utf-8') # NHWC or NCHW
        print(f'{node.name}: {layout}')

You can also do this by looking at the graph directly (graph-def variant only mentioned for the case where you already have it):

graph = func.graph
for op in graph.get_operations():
    node = op.node_def
    # do the same

Important note: sometimes (actually possibly always) the function graph might contain other function calls which might "hide" the operations you're actually interested in, so you might have to recursively delve into the sub-functions, e.g.:

def get_layouts(func, layouts=[]):
    for op in func.graph.get_operations():
        node = op.node_def
        if 'data_format' in node.attr:
            layout = node.attr['data_format'].s.decode('utf-8') # NHWC or NCHW
            layouts.append((op, layout))

    # Go through sub-functions
    subfuncs = func.graph._functions.values() # please comment if there's a better way to obtain this list than by accessing the internal-dict
    for func in subfuncs:
        get_layouts(func, layouts)
    
    return layouts

Answer 3

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: {}