When loading a saved model with multiple outputs, outputs order is not deterministic

[nitsanhasson]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No
• TensorFlow installed from (source or binary): pip
• TensorFlow version (use command below): 2.5
• Python version: 3.6.9
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: cuda 11.2
• GPU model and memory: Nvidia Quadro T1000

Describe the current behavior I have a model that outputs several tensors, each with a different shape and dtype, let's call it model A. Model's A output tensors are the input of another model B. What I would like to do is to be able to save model A (using model.save("...")), and then use it later by loading it. I found out that when loading model A and calling it, the outputs order in non deterministic, which causes a problem when I try to use those outputs as inputs to model B.

Describe the expected behavior I'd like the model to maintain the original outputs order

Standalone code to reproduce the issue Provide a reproducible test case that is the bare minimum necessary to generate the problem. If possible, please share a link to Colab/Jupyter/any notebook.

import tensorflow as tf
from tensorflow import keras

class A_Pow(keras.layers.Layer):
    def __init__(self, **kwargs):
        super(A_Pow, self).__init__(**kwargs)
        self.input_spec = tf.keras.layers.InputSpec(shape=(1, 2), dtype=tf.float32)

    def call(self, input):
        return input * input

class B_AddX(keras.layers.Layer):
    def __init__(self, x, **kwargs):
        super(B_AddX, self).__init__(**kwargs)
        self.x = x
        self.input_spec = tf.keras.layers.InputSpec(shape=(1, 3), dtype=tf.int32)

    def call(self, input):
        return input + self.x

pow_input = tf.keras.Input(shape=(2), batch_size=1, dtype=tf.float32)
addx_input = tf.keras.Input(shape=(3), batch_size=1, dtype=tf.int32)
pow_output = A_Pow()(pow_input)
addx_output = B_AddX(2)(addx_input)
model = tf.keras.Model(inputs=[pow_input, addx_input], outputs=[pow_output, addx_output])

model_output = model([tf.constant([1,2], shape=(1,2), dtype=tf.float32), tf.constant([3,4,5], shape=(1,3), dtype=tf.int32)])

model.compile()
model.save("multiple_outputs_model")

loaded_model = tf.saved_model.load("multiple_outputs_model")
loaded_model = loaded_model.signatures["serving_default"]

loaded_model_output = loaded_model(input_1=tf.constant([1,2], shape=(1,2), dtype=tf.float32), input_2=tf.constant([3,4,5], shape=(1,3), dtype=tf.int32))

print(model_output)
print(loaded_model_output)

# The order of the outputs in the loaded model is not determoinistic
# one run output:
#   [<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[1., 4.]], dtype=float32)>, <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[5, 6, 7]], dtype=int32)>]
#   {'b__add_x': <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[5, 6, 7]], dtype=int32)>, 'a__pow': <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[1., 4.]], dtype=float32)>}

# another run output:
#   [<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[1., 4.]], dtype=float32)>, <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[5, 6, 7]], dtype=int32)>]
#   {'a__pow': <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[1., 4.]], dtype=float32)>, 'b__add_x': <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[5, 6, 7]], dtype=int32)>}

# The second run has the valid output (a__pow before b_pow)

[sushreebarsa]

@nitsanhasson Could you please refer to the link and similar issue ,please let us know if it helps?Thank you!

[nitsanhasson]

@sushreebarsa Thanks for the reply, but my error is different. The outputs are correct but are in the wrong order. I thought this may be a dictionary related error since I was using Python 3.6, but the same thing happens with Python 3.9.

I need the models outputs to be in the same order both when I create and use the model without saving it, and when loading it.

Another thing I was trying to do is give the outputs names, so that when loading the saved model I can get the different tensors by their keys instead of their index, but for some reason, when loading the model those names are not retrieved.

I'm really stuck at this point because those outputs are needed for another model

[sushreebarsa]

@jvishnuvardhan Was able to reproduce the issue on Colab using TF v2.5, 2.6 and tf-nightly, Please find the gist here for reference. Thanks!

[nitsanhasson]

@jvishnuvardhan any updates?

[jvishnuvardhan]

@nitsanhasson I updated your code little bit as shown in this gist.

Mainly i changed input as a dictionary and also changed loading the keras model using tf.keras.models.load_model instead of tf.saved_model.load.

Please verify once and close the issue if this was resolved for you. Thanks!

[nitsanhasson]

@jvishnuvardhan that looks good, thanks!