Keras model pickle-able but tf.keras model not pickle-able

[Edwin-Koh1]

System information

• Windows 10
• Tensorflow 2.0 (CPU)
• joblib 0.14.0
• Python 3.7.5
• Keras 2.3.1

Hello everybody! This is my first post so please forgive me if I have missed something. So I'm trying to use a genetic algorithm to train and evaluate multiple NN architectures so I need to parallelize them on a multi-core CPU. Therefore I have used joblib to try to parallelize this. However, I was stuck on my tf.keras code because it wasn't pickleable. After many hours of debugging I finally realised that the tf.keras models are not pickleable whereas keras models are.

Describe the current behavior
The code below works but if you replaced keras with tf.keras, there will be an error:
Could not pickle the task to send it to the workers.

Describe the expected behavior
Moving forward, tf.keras should be replacing keras and therefore tf.keras should also be pickleable.

Code to reproduce the issue

#The following is a simple code to illustrate the problem:
from joblib import Parallel, delayed
import keras
import tensorflow as tf

def test():
    model = keras.models.Sequential()
    return

Parallel(n_jobs=8)(delayed(test)(i) for i in range(10)) #this works as intended

def test_tf():
    model = tf.keras.models.Sequential()
    return

Parallel(n_jobs=8)(delayed(test_tf)(i) for i in range(10)) #this will spit out the error above

Other comments
I guess a quick fix would just be to replace all the existing code with tf.keras to just keras but seeing as keras support will be discontinued and absorbed by Tensorflow 2.0, I think this should be fixed.

[ravikyram]

@Edwin-Koh1

Can you please check with nightly version(!pip install tf-nightly==2.1.0dev20191201 ) and see if the error still persists. There are lot of performance improvements in latest nightly versions. Thanks!

[ravikyram]

Automatically closing due to lack of recent activity. Please update the issue when new information becomes available, and we will reopen the issue. Thanks!

[hartikainen]

@ravikyram I'm still seeing this issue on tensorflow==2.1.0:

import pickle

import tensorflow as tf


def main():
    model_1 = tf.keras.Sequential((
        tf.keras.layers.Dense(16, activation='relu'),
        tf.keras.layers.Dense(1, activation='linear'),
    ))

    _ = model_1(tf.random.uniform((15, 3)))

    model_2 = pickle.loads(pickle.dumps(model_1))

    for w1, w2 in zip(model_1.get_weights(), model_2.get_weights()):
        tf.debugging.assert_equal(w1, w2)


if __name__ == '__main__':
    main()

results in

Traceback (most recent call last):
  File "/Users/hartikainen/conda/envs/softlearning-3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/Users/hartikainen/conda/envs/softlearning-3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/Users/hartikainen/github/rail-berkeley/softlearning-3/tests/test_pickle_keras_model.py", line 25, in <module>
    main()
  File "/Users/hartikainen/github/rail-berkeley/softlearning-3/tests/test_pickle_keras_model.py", line 18, in main
    model_2 = pickle.loads(pickle.dumps(model_1))
TypeError: can't pickle weakref objects

$ pip freeze | grep "tf\|tensor"
tensorboard==2.1.0
tensorflow==2.1.0
tensorflow-estimator==2.1.0
tensorflow-probability==0.9.0
$ python --version
Python 3.7.5

[ravikyram]

I have tried on colab with TF version 2.1.0-rc2, 2.2.0-dev20200113 and was able to reproduce the issue.Please, find the gist here. Thanks!

[hartikainen]

@ravikyram, should keras functional models be picklable too or not? I'd assume if Sequential models are then functional models should too be? Or does functional models have some properties that make them harder to pickle?

$ python -m tests.test_pickle_keras_functional_model
2020-01-17 16:47:08.567598: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-01-17 16:47:08.581327: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fa0a55aa6c0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-01-17 16:47:08.581362: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
Traceback (most recent call last):
  File "/Users/hartikainen/conda/envs/softlearning-3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/Users/hartikainen/conda/envs/softlearning-3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/Users/hartikainen/github/rail-berkeley/softlearning-3/tests/test_pickle_keras_functional_model.py", line 20, in <module>
    main()
  File "/Users/hartikainen/github/rail-berkeley/softlearning-3/tests/test_pickle_keras_functional_model.py", line 13, in main
    model_2 = pickle.loads(pickle.dumps(model_1))
TypeError: can't pickle _thread.RLock objects

[marfox]

Hi everyone,
I'm trying to switch from standalone keras to tensorflow.keras as per the recommendation at https://keras.io/.
I'm hitting the same exception as #34697 (comment) with joblib (which uses pickle under the hood).

System information:

• Debian 10 (buster)
• Python 3.7.6
• joblib 0.14.1
• tensorflow 2.1.0

Script to reproduce:

import joblib
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

model = Sequential()
model.add(Dense(1, input_dim=42, activation='sigmoid'))
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
joblib.dump(model, 'model.pkl')

Output:

TypeError: can't pickle _thread.RLock objects

[epetrovski]

Here's a fix adapted from http://zachmoshe.com/2017/04/03/pickling-keras-models.html intended for solving the same issue back when Keras models used to not be pickleable.

import pickle
import tempfile
from tensorflow.keras.models import Sequential, load_model, save_model, Model
from tensorflow.keras.layers import Dense

# Hotfix function
def make_keras_picklable():
    def __getstate__(self):
        model_str = ""
        with tempfile.NamedTemporaryFile(suffix='.hdf5', delete=True) as fd:
            save_model(self, fd.name, overwrite=True)
            model_str = fd.read()
        d = {'model_str': model_str}
        return d

    def __setstate__(self, state):
        with tempfile.NamedTemporaryFile(suffix='.hdf5', delete=True) as fd:
            fd.write(state['model_str'])
            fd.flush()
            model = load_model(fd.name)
        self.__dict__ = model.__dict__


    cls = Model
    cls.__getstate__ = __getstate__
    cls.__setstate__ = __setstate__

# Run the function
make_keras_picklable()

# Create the model
model = Sequential()
model.add(Dense(1, input_dim=42, activation='sigmoid'))
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])

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

[leandrocouto]

@epetrovski Should I call this code whenever I'm about to pickle a model or can I just call it at the beginning of my application (before creating the model)?

[epetrovski]

@epetrovski Should I call this code whenever I'm about to pickle a model or can I just call it at the beginning of my application (before creating the model)?

You can definitely just call it once at the beginning of your app after importing tensorflow.keras.models.Model. Executing the function adds two new methods __getstate__()and __setstate__() to the tensorflow.keras.models.Model class so it should work every time you want to pickle a member of the updated tf.keras Model class - ie. your own model.

[adriangb]

Here is an alternative to @epetrovski 's answer that does not require saving to a file:

import pickle

from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense
from tensorflow.python.keras.layers import deserialize, serialize
from tensorflow.python.keras.saving import saving_utils


def unpack(model, training_config, weights):
    restored_model = deserialize(model)
    if training_config is not None:
        restored_model.compile(
            **saving_utils.compile_args_from_training_config(
                training_config
            )
        )
    restored_model.set_weights(weights)
    return restored_model

# Hotfix function
def make_keras_picklable():

    def __reduce__(self):
        model_metadata = saving_utils.model_metadata(self)
        training_config = model_metadata.get("training_config", None)
        model = serialize(self)
        weights = self.get_weights()
        return (unpack, (model, training_config, weights))

    cls = Model
    cls.__reduce__ = __reduce__

# Run the function
make_keras_picklable()

# Create the model
model = Sequential()
model.add(Dense(1, input_dim=42, activation='sigmoid'))
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])

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

Source: https://docs.python.org/3/library/pickle.html#object.__reduce__

I feel like maybe this could be added to Model? Are there any cases where this would not work?

[sushreebarsa]

@Edwin-Koh1 Could you please refer to the link1 ,link2 and let us know if it helps ? Thanks!

[Edwin-Koh1]

Hey all, I believe the issue was resolved by @adriangb in #39609 and the keras pull #14748. Should be fixed in Keras and Tensorflow 2.6.0 but have not tested this yet.

[ageron]

@Edwin-Koh1 , indeed, the problem seems fixed in TF 2.6.0! 👍

This now works (while it failed in TF 2.5.1):

import joblib
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

model = Sequential([Dense(1, input_shape=[42], activation='sigmoid')])
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
joblib.dump(model, 'model.pkl')

[adriangb]

Glad we could fix it for you @ageron! It should also work with Functional models by the way.

@Edwin-Koh1 should we close the issue?

[ageron]

It works with joblib in TF 2.6.0, but not with pickle, however:

import pickle
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

model = Sequential([Dense(1, input_shape=[42], activation='sigmoid')])
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
with open('model.pkl', 'wb') as f:
  pickle.dump(model, f)

Raises:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-e391f4844d65> in <module>()
      6 model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
      7 with open('model.pkl', 'wb') as f:
----> 8   pickle.dump(model, f)

TypeError: can't pickle weakref objects

[adriangb]

Can you try with tf-nightly? I'm wondering if the fix was actually not in 2.6.0 and your example just happens to be broken in 2.5.1. In this colab notebook, on tf-nightly, your example runs fine.

[sushreebarsa]

@Edwin-Koh1 Could you please let us know if this issue is resolved for you ? If it is resolved please feel free to move this ticket to closed status .Thank you!

[ageron]

@adriangb you're right, both the joblib and pickle code examples work with tf-nightly (while only the joblib example worked in 2.6.0). 👍

[google-ml-butler]

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]

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

[google-ml-butler]

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

[raj-gupta1]

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.

Issue resolved

[AhmedCoolProjects]

Here's a fix adapted from http://zachmoshe.com/2017/04/03/pickling-keras-models.html intended for solving the same issue back when Keras models used to not be pickleable.

import pickle
import tempfile
from tensorflow.keras.models import Sequential, load_model, save_model, Model
from tensorflow.keras.layers import Dense

# Hotfix function
def make_keras_picklable():
    def __getstate__(self):
        model_str = ""
        with tempfile.NamedTemporaryFile(suffix='.hdf5', delete=True) as fd:
            save_model(self, fd.name, overwrite=True)
            model_str = fd.read()
        d = {'model_str': model_str}
        return d

    def __setstate__(self, state):
        with tempfile.NamedTemporaryFile(suffix='.hdf5', delete=True) as fd:
            fd.write(state['model_str'])
            fd.flush()
            model = load_model(fd.name)
        self.__dict__ = model.__dict__


    cls = Model
    cls.__getstate__ = __getstate__
    cls.__setstate__ = __setstate__

# Run the function
make_keras_picklable()

# Create the model
model = Sequential()
model.add(Dense(1, input_dim=42, activation='sigmoid'))
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])

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

That really worked for me, But I got an issue while using pickle to load the model, so I should to keep using tensorflow so the pickle load function can load the model.

[raghunathanp95]

It works with joblib in TF 2.6.0, but not with pickle, however:

import pickle
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

model = Sequential([Dense(1, input_shape=[42], activation='sigmoid')])
model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
with open('model.pkl', 'wb') as f:
  pickle.dump(model, f)

Raises:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-e391f4844d65> in <module>()
      6 model.compile(optimizer='Nadam', loss='binary_crossentropy', metrics=['accuracy'])
      7 with open('model.pkl', 'wb') as f:
----> 8   pickle.dump(model, f)

TypeError: can't pickle weakref objects

import tensorflow as tf
tf.keras.models.save_model(model, filename)

Try this instead of pickle.dump or joblilb.dump