keras icon indicating copy to clipboard operation
keras copied to clipboard

Add registered custom objects inside pickled model file

Open mthiboust opened this issue 2 months ago • 3 comments

See the discussion in https://github.com/keras-team/keras/issues/19832

Context

Currently, keras serialization does not store the definition of custom objects (see here). Thus, pickled models are not self-sufficient if you want to load them in a separate vanilla session and call model.predict() if they contain such objects.

In practice, if you have 2 isolated codebases for training and inference, you always need to update the inference one along with the training one when using a custom object. It does not allow to experiment quickly with new models if your inference codebase has long CI/CD pipelines. Storing the definition inside the pickle file would allow to decouple training from inference.

Suggested change

Modify the __reduce__() method of KerasSaveable to store the definition of registered custom keras objects inside the pickle file when a keras model is serialized via pickle-like libraries (e.g. pickle, cloudpickle, dill). Registered custom keras objects are pickled with the standard pickle to avoid extra dependencies (but inheriting pickle limitations like lambda functions, this is a tradeoff).

Add a simple test with a registered custom layer. This test would ideally load the pickled model in a new subprocess to completely verify the wanted behavior. But it would complexify the test. Or maybe there is another way to test this? I can investigate this point if you want.

mthiboust avatar Jun 17 '24 19:06 mthiboust