RFC: Pickle Protocol for Keras

[adriangb]

Comment period is open until 10-13.

As per discussion in tensorflow/tensorflow#39609 (comment)

Status	Proposed
RFC #	286
Author(s)	Adrian Garcia Badaracco ({firstname}at{firstname}gb.com), Scott Sievert (tf-rfc@stsievert.com)
Sponsor	Mihai Maruseac (mihaimaruseac@google.com)
Updated	2020-09-21

Objective

Implement support for Pickle, Python's serialization protocol within Keras.

[adriangb]

@stsievert

[adriangb]

resolved

[googlebot]

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[adriangb]

cc @mihaimaruseac for initial review

[mihaimaruseac]

It seems there are no more comments. Let's schedule an internal review then. I'll come back with the date.

[mihaimaruseac]

Design review scheduled for 22nd of October at 10am Pacific Time.

[adriangb]

Design review scheduled for 22nd of October at 10am Pacific Time.

Awesome, thank you!

[fchollet]

What is the list of base classes that you expect will need to implement the __reduce__ method?

[adriangb]

I think a good starting point would be:

• tf.keras.Model
• tf.keras.metrics.Metric
• tf.keras.losses.Loss
• tf.keras.optimizers.Optimizer
• tf.keras.callbacks.Callback

I think layers are usually serialized as part of a Model; I don't think there's much of a use case for implementing __reduce__ for individual layers.

But generally, all classes that already have a serialize and deserialize method could benefit getting a __reduce__ that just references those, but there is no reason why they have all be implemented together (or at all).

[ematejska]

Design review scheduled for 22nd of October at 10am Pacific Time.

That's at 3-4PM Pacific time. Sorry for any confustion.

[mihaimaruseac]

Notes from the design meeting:

1. Only one question: Base classes that need to implement __reduce__:
   • Do we need the Layer class to be listed?
     • we're pickling as part of the model
     • Francois: if part of a model, it makes sense to also be available on the layer
2. Francois: High level API: this is very useful for users, good proposal
   • makes sense for models, not so much for optimizers (since state needs to be maintained, we only serialize configs)
   • if all we need is the config, get_config() works, if we need state then we need more
     • add some intermediate SavedModel serialization
     • or, current RFC only targets model case, is very simple, can be implemented and we'll make a new one for the other case
       • we should not implement at the moment for objects (callbacks, optimizers, metrics, losses with lambdas, etc.) where we only promise to save the config, not the internal state (as this breaks guarantees and API safety; state is usually retrieved via different API calls; results of these can be used to set state too)
3. Callbacks need to be serializable before we can implement this fully
4. RFC only covers built-in objects; in many cases models have custom layers => will the proposal still work?
   • TF function is not something that can be pickled
   • supposed to work with everything that is serializable via SavedModel
     • serialization via SavedModel which can save to RAM
     • why not implement this on the SavedModel? It is used outside in other packages (i.e., scikeras)
5. If we pickle something that we won't support, will we throw an error?
   • with the model, it should work
6. Does this impose new versioning requirements?
   • since it is backed by SavedModel, no, just uses SavedModel's versioning
   • pickle itself makes no guarantee about versioning, so this is more strict

Conclusions:

1. We will restrict scope to SavedModel pickling as this solves majority of user requirements / feature requests
2. Proposal accepted

[jakirkham]

Great to see this in! Thanks all! 😄 👏