keras
keras copied to clipboard
Add registered custom objects inside pickled model file
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.