Custom objects support when pickling keras models

[mthiboust]

Following the recent work on pickle support for Keras 3 (PRs https://github.com/keras-team/keras/pull/19555 & https://github.com/keras-team/keras/pull/19592), it would be useful to add the definition of custom objects directly inside the pickled file. In practice, it allows to decouple the training from the inference with a generic inference code (no need to have the custom object definition available in the environment where the model is loaded).

For now, I am toying with this monkey patch:

import io

import cloudpickle as pickle
from keras.utils import get_custom_objects
from keras.src.saving.keras_saveable import KerasSaveable


# Monkey patch the KerasSaveable class to save custom keras objects alongwith the model


def new_unpickle_model(cls, model_buf, custom_objects_buf):
    import keras.src.saving.saving_lib as saving_lib

    # pickle is not safe regardless of what you do.
    custom_objects = pickle.load(custom_objects_buf)
    return saving_lib._load_model_from_fileobj(
        model_buf, custom_objects=custom_objects, compile=True, safe_mode=False
    )


def new__reduce__(self):
    """__reduce__ is used to customize the behavior of `pickle.pickle()`.

    The method returns a tuple of two elements: a function, and a list of
    arguments to pass to that function.  In this case we just leverage the
    keras saving library."""
    import keras.src.saving.saving_lib as saving_lib

    model_buf = io.BytesIO()
    saving_lib._save_model_to_fileobj(self, model_buf, "h5")

    custom_objects_buf = io.BytesIO()
    pickle.dump(get_custom_objects(), custom_objects_buf)
    custom_objects_buf.seek(0)

    return (
        self._unpickle_model,
        (model_buf, custom_objects_buf),
    )


KerasSaveable._unpickle_model = classmethod(new_unpickle_model)
KerasSaveable.__reduce__ = new__reduce__

Is it the right way to do it? Would you welcome such a change to be directly integrated into https://github.com/keras-team/keras/blob/master/keras/src/saving/keras_saveable.py ?

I know that pickle is not recommended for security reasons when you load a pickled model from an untrusted source. But there are cases where the same entity both dumps and loads the models.

Tagging @LukeWood who may have an opinion on this.