Keras custom layers and SavedModel export: which pieces survive the round trip without custom_objects?
0 reputation · 27 Mar 2021, 07:26 UTC
0 reputation · 27 Mar 2021, 07:26 UTC
I am designing a deployment pipeline where a tf.keras model containing several custom layers must be exported via model.save() in SavedModel format and later reloaded in a separate process, potentially by a team that does not have the original layer classes importable. Assume TensorFlow 2.x with eager execution.
The documented contract is that custom layers implement get_config() (and optionally from_config()) so the model config can reconstruct them, and that tf.keras.models.load_model() needs a custom_objects mapping when the classes are not registered. What is unclear to me is the boundary between what the SavedModel format itself preserves (the traced computation graph and weights, loadable via tf.saved_model.load() with no Python class definitions) and what the Keras-level reload strictly requires.
This matters because the consuming side may only need inference signatures, while other consumers need a fully functional Keras model for continued training.
Specifically:
tf.saved_model.load() signatures even when the custom layer classes are unavailable?@tf.keras.utils.register_keras_serializable() sufficient to drop the custom_objects argument on reload, and does that registration survive across processes?get_config() alone is insufficient, e.g. layers with non-serializable constructor arguments?A thoughtful contribution can make all the difference. Be the first to share one.
Use comments to ask for clarification. Post a solution as an answer.
No question comments on this page.