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.