tf.keras.layers.Layer.weights has inconsistent order in 'saved_model' format

[MathiasToftas]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): yes
• OS Platform and Distribution: Linux Ubuntu 20.04
• TensorFlow installed from: binary
• TensorFlow version: v2.9.0-18-gd8ce9f9c301 2.9.1
• Python version: 3.8
• Exact command to reproduce: see google colab page

Describe the problem.

I originaly posted this issue in the tensorflow github, but was asked to move it here (original issue: https://github.com/tensorflow/tensorflow/issues/57438)

I'm building tensorflow lite and tensorflow 'saved_model's with the purpose of training them on edge devices in a federated learning setting. To start off I've been following this guide (On-Device Training with TensorFlow Lite).

To perform the FL, the weights of each edge model must be transferred to the server. The way I've done this is by accessing the weights using the tf.keras.layers.Layer.weights member. The issue lies in how the weights of this list are not always ordered consistently. If these weights are not ordered, the server will need some kind of extra information to sort the weights, which needlessly complicates the weights transfer.

Describe the current behavior.

Currently, if you access tf.keras.layers.Layer.weights from a normal keras model or tf.Module that exists in memory (no saving, loading, compressing, etc) the weights are given as a list in an order that is consistent, i.e. we can reinitialize the model and still receive the weights in the same order. This is the expected behavior.

However, if you save the same model in a 'saved_model' format (which is the recommended way to construct a tensorflow lite model Convert TensorFlow models), the order in which these weights are given is seemingly randomized each time the model is loaded.

Confusingly enough, if we convert the 'saved_model' into a tensorflow light model, the order is once again consistent as expected.

Describe the expected behavior.

The expected behavior is that accessing tf.keras.layers.Layer.weights should always produce the exact same result, no matter how the model has been packaged. Especially given that for the above 3 cases, a single model source file is used.

Contributing.

• Do you want to contribute a PR? (yes/no): no
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue.

colab

[sushreebarsa]

@gadagashwini I was able to replicate the issue on colab, please find the gist here. Thank you!

[gadagashwini]

HI @MathiasToftas, You need to use signatures under tf.function ( for graph mode). Please take a look at this reference link for more information. Thank you!

[MathiasToftas]

I'm not sure I follow, from the link you've shared, and from my own experience, you only need to add the 'input' signature under @tf.function. The function show in the colab doc does not have any inputs (other than self), so I'm not sure what can be added.

From what the link you shared seems to say, the outputs signature is gathered from the structure of the dict you return, and looking at the returned dict in the colab doc, there is nothing there that would create a random order, except for the Layer.weights attribute.

Further, if this was the issue, why would the conversion to a tf lite model even succeed, and why would the signature_runner return anything at all if the graph conversion was unsuccessful?

[google-ml-butler[bot]]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler[bot]]

Closing as stale. Please reopen if you'd like to work on this further.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No