212 lines
6.9 KiB
Python
212 lines
6.9 KiB
Python
import inspect
|
|
|
|
from keras_core.api_export import keras_core_export
|
|
from keras_core.backend.common import global_state
|
|
|
|
GLOBAL_CUSTOM_OBJECTS = {}
|
|
GLOBAL_CUSTOM_NAMES = {}
|
|
|
|
|
|
@keras_core_export(
|
|
[
|
|
"keras_core.saving.CustomObjectScope",
|
|
"keras_core.saving.custom_object_scope",
|
|
]
|
|
)
|
|
class CustomObjectScope:
|
|
"""Exposes custom classes/functions to Keras deserialization internals.
|
|
|
|
Under a scope `with custom_object_scope(objects_dict)`, Keras methods such
|
|
as `keras_core.models.load_model()` or
|
|
`keras_core.models.model_from_config()` will be able to deserialize any
|
|
custom object referenced by a saved config (e.g. a custom layer or metric).
|
|
|
|
Example:
|
|
|
|
Consider a custom regularizer `my_regularizer`:
|
|
|
|
```python
|
|
layer = Dense(3, kernel_regularizer=my_regularizer)
|
|
# Config contains a reference to `my_regularizer`
|
|
config = layer.get_config()
|
|
...
|
|
# Later:
|
|
with custom_object_scope({'my_regularizer': my_regularizer}):
|
|
layer = Dense.from_config(config)
|
|
```
|
|
|
|
Args:
|
|
custom_objects: Dictionary of `{name: object}` pairs.
|
|
"""
|
|
|
|
def __init__(self, custom_objects):
|
|
self.custom_objects = custom_objects or {}
|
|
self.backup = None
|
|
|
|
def __enter__(self):
|
|
self.backup = global_state.get_global_attribute(
|
|
"custom_objects_scope_dict", {}
|
|
).copy()
|
|
global_state.set_global_attribute(
|
|
"custom_objects_scope_dict", self.custom_objects.copy()
|
|
)
|
|
return self
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
global_state.set_global_attribute(
|
|
"custom_objects_scope_dict", self.backup.copy()
|
|
)
|
|
|
|
|
|
# Alias.
|
|
custom_object_scope = CustomObjectScope
|
|
|
|
|
|
@keras_core_export(
|
|
"keras_core.saving.get_custom_objects",
|
|
)
|
|
def get_custom_objects():
|
|
"""Retrieves a live reference to the global dictionary of custom objects.
|
|
|
|
Custom objects set using using `custom_object_scope()` are not added to the
|
|
global dictionary of custom objects, and will not appear in the returned
|
|
dictionary.
|
|
|
|
Example:
|
|
|
|
```python
|
|
get_custom_objects().clear()
|
|
get_custom_objects()['MyObject'] = MyObject
|
|
```
|
|
|
|
Returns:
|
|
Global dictionary mapping registered class names to classes.
|
|
"""
|
|
return GLOBAL_CUSTOM_OBJECTS
|
|
|
|
|
|
@keras_core_export(
|
|
"keras_core.saving.register_keras_serializable",
|
|
)
|
|
def register_keras_serializable(package="Custom", name=None):
|
|
"""Registers an object with the Keras serialization framework.
|
|
|
|
This decorator injects the decorated class or function into the Keras custom
|
|
object dictionary, so that it can be serialized and deserialized without
|
|
needing an entry in the user-provided custom object dict. It also injects a
|
|
function that Keras will call to get the object's serializable string key.
|
|
|
|
Note that to be serialized and deserialized, classes must implement the
|
|
`get_config()` method. Functions do not have this requirement.
|
|
|
|
The object will be registered under the key `'package>name'` where `name`,
|
|
defaults to the object name if not passed.
|
|
|
|
Example:
|
|
|
|
```python
|
|
# Note that `'my_package'` is used as the `package` argument here, and since
|
|
# the `name` argument is not provided, `'MyDense'` is used as the `name`.
|
|
@register_keras_serializable('my_package')
|
|
class MyDense(keras_core.layers.Dense):
|
|
pass
|
|
|
|
assert get_registered_object('my_package>MyDense') == MyDense
|
|
assert get_registered_name(MyDense) == 'my_package>MyDense'
|
|
```
|
|
|
|
Args:
|
|
package: The package that this class belongs to. This is used for the
|
|
`key` (which is `"package>name"`) to idenfify the class. Note that
|
|
this is the first argument passed into the decorator.
|
|
name: The name to serialize this class under in this package. If not
|
|
provided or `None`, the class' name will be used (note that this is
|
|
the case when the decorator is used with only one argument, which
|
|
becomes the `package`).
|
|
|
|
Returns:
|
|
A decorator that registers the decorated class with the passed names.
|
|
"""
|
|
|
|
def decorator(arg):
|
|
"""Registers a class with the Keras serialization framework."""
|
|
class_name = name if name is not None else arg.__name__
|
|
registered_name = package + ">" + class_name
|
|
|
|
if inspect.isclass(arg) and not hasattr(arg, "get_config"):
|
|
raise ValueError(
|
|
"Cannot register a class that does not have a "
|
|
"get_config() method."
|
|
)
|
|
|
|
GLOBAL_CUSTOM_OBJECTS[registered_name] = arg
|
|
GLOBAL_CUSTOM_NAMES[arg] = registered_name
|
|
|
|
return arg
|
|
|
|
return decorator
|
|
|
|
|
|
@keras_core_export("keras_core.saving.get_registered_name")
|
|
def get_registered_name(obj):
|
|
"""Returns the name registered to an object within the Keras framework.
|
|
|
|
This function is part of the Keras serialization and deserialization
|
|
framework. It maps objects to the string names associated with those objects
|
|
for serialization/deserialization.
|
|
|
|
Args:
|
|
obj: The object to look up.
|
|
|
|
Returns:
|
|
The name associated with the object, or the default Python name if the
|
|
object is not registered.
|
|
"""
|
|
if obj in GLOBAL_CUSTOM_NAMES:
|
|
return GLOBAL_CUSTOM_NAMES[obj]
|
|
else:
|
|
return obj.__name__
|
|
|
|
|
|
@keras_core_export("keras_core.saving.get_registered_object")
|
|
def get_registered_object(name, custom_objects=None, module_objects=None):
|
|
"""Returns the class associated with `name` if it is registered with Keras.
|
|
|
|
This function is part of the Keras serialization and deserialization
|
|
framework. It maps strings to the objects associated with them for
|
|
serialization/deserialization.
|
|
|
|
Example:
|
|
|
|
```python
|
|
def from_config(cls, config, custom_objects=None):
|
|
if 'my_custom_object_name' in config:
|
|
config['hidden_cls'] = tf.keras.saving.get_registered_object(
|
|
config['my_custom_object_name'], custom_objects=custom_objects)
|
|
```
|
|
|
|
Args:
|
|
name: The name to look up.
|
|
custom_objects: A dictionary of custom objects to look the name up in.
|
|
Generally, custom_objects is provided by the user.
|
|
module_objects: A dictionary of custom objects to look the name up in.
|
|
Generally, module_objects is provided by midlevel library
|
|
implementers.
|
|
|
|
Returns:
|
|
An instantiable class associated with `name`, or `None` if no such class
|
|
exists.
|
|
"""
|
|
custom_objects_scope_dict = global_state.get_global_attribute(
|
|
"custom_objects_scope_dict", {}
|
|
)
|
|
if name in custom_objects_scope_dict:
|
|
return custom_objects_scope_dict[name]
|
|
elif name in GLOBAL_CUSTOM_OBJECTS:
|
|
return GLOBAL_CUSTOM_OBJECTS[name]
|
|
elif custom_objects and name in custom_objects:
|
|
return custom_objects[name]
|
|
elif module_objects and name in module_objects:
|
|
return module_objects[name]
|
|
return None
|