Cannot deserialize Adam optimizer on M1/M2 mac

[tomcarter23]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Mac arm M2 Ventura 13.3
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.13.0
• Python version: 3.11.4
• Bazel version (if compiling from source):
• GPU model and memory: Apple M2 Pro 16 GB
• Exact command to reproduce:

Describe the problem.

Cannot deserialize a model compiled with the tf.keras.optimizers.Adam optimizer on M1/M2 macs.

Describe the current behavior.

When creating a Keras model on a M1/M2 mac the following messages are displayed indicating that the default optimizer tf.keras.optimizers.Adam runs slowly on M1/M2 macs. Keras then "falls back" to the legacy optimizer tf.keras.optimizers.legacy.Adam.

WARNING:absl:At this time, the v2.11+ optimizer `tf.keras.optimizers.Adam` runs slowly on M1/M2 Macs, please use the legacy Keras optimizer instead, located at `tf.keras.optimizers.legacy.Adam`.
WARNING:absl:There is a known slowdown when using v2.11+ Keras optimizers on M1/M2 Macs. Falling back to the legacy Keras optimizer, i.e., `tf.keras.optimizers.legacy.Adam`.

Then, when serializing and trying to deserialize the model the following error is raised:

self = <keras.src.optimizers.legacy.adam.Adam object at 0x29b503a90>
name = 'build'

    def __getattribute__(self, name):
        """Overridden to support hyperparameter access."""
        try:
>           return super().__getattribute__(name)
E           AttributeError: 'Adam' object has no attribute 'build'

Indicating that the legacy optimizer cannot be deserialized because it does not implement the build attribute. On all other platforms the model can be serialized and deserialized because the tf.keras.optimizers.Adam is used rather than the legacy version.

Describe the expected behavior.

The expected behaviour is that the model can be deserialized without raising an attribute error.

Standalone code to reproduce the issue.

import tensorflow as tf

import pickle


model = tf.keras.Model(inputs={"a": tf.keras.Input(shape=(10,))}, outputs={"b": tf.keras.Input(shape=(10,))})
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='mse')

with open('test.pkl', 'wb') as f:
    pickle.dump(model, f)

with open('test.pkl', 'rb') as f:
    opt = pickle.load(f)

[tilakrayal]

@sachinprasadhs,
I was able to execute the mentioned code without any issues on tensorflow v2.12, v2.13 but whereas on Mac M1 the code was failing with the above mentioned error. Kindly find the gist and reference below.

[sachinprasadhs]

I tried to replicate the behavior with keras_core, but I'm getting different error both in Colab and M1.
Here is the gist attached for reference. https://gist.github.com/sachinprasadhs/58079b64a4beb47d547ac1529f25745b

[haifeng-jin]

Temporarily assigned to @nkovela1.
Please reassign if it is not relevant.

[psfoley]

As an additional datapoint, the root issue here is wider spread than M1/M2 CPUs. Any attempted serialization and deserialization of legacy optimizers results in the same failure:

Standalone code to reproduce the issue.

import tensorflow as tf

import pickle


model = tf.keras.Model(inputs={"a": tf.keras.Input(shape=(10,))}, outputs={"b": tf.keras.Input(shape=(10,))})
model.compile(optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=0.001), loss='mse')

with open('test.pkl', 'wb') as f:
   pickle.dump(model, f)

with open('test.pkl', 'rb') as f:
   opt = pickle.load(f)

[nkovela1]

Hi @tomcarter23 , we actively discourage pickling for serialization/deserialization. Please try using themodel.to_json() and keras.models.model_from_json() methods.

[tomcarter23]

Hi,

I want to be able to pickle the whole model so I can ship it across a spark cluster and keep information like the training history and the optimizer state. As far as I know this is not possible using the to_json() and from_model_json() but please correct me if I am wrong.

On every platform except for M1/M2 mac I can pickle the model and ship it across spark nodes. So currently I have to have an explicit check for a tf.keras.optimizers.legacy.Optimizer instance (M1/M2 mac) and deal with the shipping of the model in a different and less optimal way, that doesn't preserve the training history and the optimizer state.

[tomcarter23]

Hi, wondering what is happening with this

[nkovela1]

Hi @tomcarter23, sorry for the delay. I have a pending fix, but Keras is currently undergoing an internal PR freeze in preparation for Keras to become multi-backend. I will keep you updated once I'm able to submit a fix.

[tomcarter23]

Thats great! Thanks :)

[tomcarter23]

Hi, is there any update here?

[nkovela1]

Hi Tom, yes! A fix has been pushed and is already available on TF-nightly. It should be part of the Keras 2.15 release.

Thank you again for your patience!

[tomcarter23]

Great, thanks for that! could you link the PR please?

[nkovela1]

Sure! This is the PR: keras-team/keras#18492

[tilakrayal]

@tomcarter23,
I tried to execute the mentioned code on Apple M1 with the tf-nightly(2.15.0-dev20230927) and it was executed without any fail or error. Kindly find the screenshot below for the reference.

Thank you!