EVOLUTION-MANAGER

Edit File: serialized_attributes.py

# Copyright 2018 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Helper classes that list&validate all attributes to serialize to SavedModel. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.python.eager import def_function from tensorflow.python.eager import function as defun from tensorflow.python.keras.saving.saved_model import constants from tensorflow.python.keras.utils.generic_utils import LazyLoader from tensorflow.python.training.tracking import base as trackable from tensorflow.python.training.tracking.tracking import AutoTrackable # TODO(b/134426265): Switch back to single-quotes to match the rest of the file # once the issue with copybara is fixed. # pylint:disable=g-inconsistent-quotes base_layer = LazyLoader( "base_layer", globals(), "tensorflow.python.keras.engine.base_layer") training_lib = LazyLoader( "training_lib", globals(), "tensorflow.python.keras.engine.training") metrics = LazyLoader("metrics", globals(), "tensorflow.python.keras.metrics") recurrent = LazyLoader( "recurrent", globals(), "tensorflow.python.keras.layers.recurrent") # pylint:enable=g-inconsistent-quotes class SerializedAttributes(object): """Class that tracks and validates all serialization attributes. Keras models contain many Python-defined components. For example, the trainable_variable property lists the model's trainable variables by recursively retrieving the trainable variables from each of the child layers. Another example is model.call, a python function that calls child layers and adds ops to the backend graph. Only Tensorflow checkpointable objects and functions can be serialized to SavedModel. Serializing a Keras model as-is results in a checkpointable object that does not resemble a Keras model at all. Thus, extra checkpointable objects and functions must be created during serialization. **Defining new serialized attributes** Child classes should be defined using: SerializedAttributes.with_attributes( 'name', checkpointable_objects=[...], functions=[...], copy_from=[...]) This class is used to cache generated checkpointable objects and functions, ensuring that new objects and functions are generated a single time. **Usage during serialization** Each Layer/Model object should have a corresponding instance of SerializedAttributes. Create a new instance by calling `SerializedAttributes.new(obj)`. Objects and functions may be saved using `.set_and_validate_checkpointable_objects`/`.set_and_and_validate_functions`. The properties `.checkpointable_objects` and `.functions` returns the cached values. **Adding/changing attributes to save to SavedModel** 1. Change the call to `SerializedAttributes.with_attributes` in the correct class: - CommonEndpoints: Base attributes to be added during serialization. If these attributes are present in a Trackable object, it can be deserialized to a Keras Model. - LayerAttributes: Attributes to serialize for Layer objects. - ModelAttributes: Attributes to serialize for Model objects. 2. Update class docstring 3. Update arguments to any calls to `set_and_validate_*`. For example, if `call_raw_tensors` is added to the ModelAttributes function list, then a `call_raw_tensors` function should be passed to `set_and_validate_functions`. **Common endpoints vs other attributes** Only common endpoints are attached directly to the root object. Keras-specific attributes are saved to a separate trackable object with the name "keras_api". The number of objects attached to the root is limited because any naming conflicts will cause user code to break. Another reason is that this will only affect users who call `tf.saved_model.load` instead of `tf.keras.models.load_model`. These are advanced users who are likely to have defined their own tf.functions and trackable objects. The added Keras-specific attributes are kept out of the way in the "keras_api" namespace. Properties defined in this class may be used to filter out keras-specific attributes: - `functions_to_serialize`: Returns dict of functions to attach to the root object. - `checkpointable_objects_to_serialize`: Returns dict of objects to attach to the root object (including separate trackable object containing keras-specific attributes) All changes to the serialized attributes must be backwards-compatible, so attributes should not be removed or modified without sufficient justification. """ @staticmethod def with_attributes( name, checkpointable_objects=None, functions=None, copy_from=None): """Creates a subclass with all attributes as specified in the arguments. Args: name: Name of subclass checkpointable_objects: List of checkpointable objects to be serialized in the SavedModel. functions: List of functions to be serialized in the SavedModel. copy_from: List of other SerializedAttributes subclasses. The returned class will copy checkpoint objects/functions from each subclass. Returns: Child class with attributes as defined in the `checkpointable_objects` and `functions` lists. """ checkpointable_objects = checkpointable_objects or [] functions = functions or [] if copy_from is not None: for cls in copy_from: checkpointable_objects.extend(cls.all_checkpointable_objects) functions.extend(cls.all_functions) classdict = { 'all_checkpointable_objects': set(checkpointable_objects), 'all_functions': set(functions)} return type(name, (SerializedAttributes,), classdict) @staticmethod def new(obj): """Returns a new SerializedAttribute object.""" if isinstance(obj, training_lib.Model): return ModelAttributes() elif isinstance(obj, metrics.Metric): return MetricAttributes() elif isinstance(obj, recurrent.RNN): return RNNAttributes() elif isinstance(obj, base_layer.Layer): return LayerAttributes() else: raise TypeError('Internal error during serialization: Expected Keras ' 'Layer object, got {} of type {}'.format(obj, type(obj))) def __init__(self): self._object_dict = {} self._function_dict = {} self._keras_trackable = AutoTrackable() @property def functions(self): """Returns dictionary of all functions.""" return {key: value for key, value in self._function_dict.items() if value is not None} @property def checkpointable_objects(self): """Returns dictionary of all checkpointable objects.""" return {key: value for key, value in self._object_dict.items() if value is not None} @property def functions_to_serialize(self): """Returns functions to attach to the root object during serialization.""" return {key: value for key, value in self.functions.items() if key in CommonEndpoints.all_functions} @property def objects_to_serialize(self): """Returns objects to attach to the root object during serialization.""" objects = {key: value for key, value in self.checkpointable_objects.items() if key in CommonEndpoints.all_checkpointable_objects} objects[constants.KERAS_ATTR] = self._keras_trackable return objects def set_and_validate_functions(self, function_dict): """Saves function dictionary, and validates dictionary values.""" for key in self.all_functions: if key in function_dict: if (function_dict[key] is not None and # Not all functions are required not isinstance(function_dict[key], (defun.Function, def_function.Function))): raise ValueError( 'Function dictionary contained a non-function object: {} (for key' ' {})'.format(function_dict[key], key)) self._function_dict[key] = function_dict[key] setattr(self._keras_trackable, key, function_dict[key]) else: raise ValueError('Function {} missing from serialized function dict.' .format(key)) return self.functions def set_and_validate_objects(self, object_dict): """Saves objects to a dictionary, and validates the values.""" for key in self.all_checkpointable_objects: if key in object_dict: if not isinstance(object_dict[key], trackable.Trackable): raise ValueError( 'Object dictionary contained a non-trackable object: {} (for key' ' {})'.format(object_dict[key], key)) self._object_dict[key] = object_dict[key] setattr(self._keras_trackable, key, object_dict[key]) else: raise ValueError( 'Object {} missing from serialized object dict.'.format(key)) return self.checkpointable_objects class CommonEndpoints(SerializedAttributes.with_attributes( 'CommonEndpoints', checkpointable_objects=['variables', 'trainable_variables', 'regularization_losses'], functions=['__call__', 'call_and_return_all_conditional_losses', '_default_save_signature'])): """Common endpoints shared by all models loadable by Keras. List of all attributes: variables: List of all variables in the model and its sublayers. trainable_variables: List of all trainable variables in the model and its sublayers. regularization_losses: List of all unconditional losses (losses not dependent on the inputs) in the model and its sublayers. __call__: Function that takes inputs and returns the outputs of the model call function. call_and_return_all_conditional_losses: Function that returns a tuple of (call function outputs, list of all losses that depend on the inputs). _default_save_signature: Traced model call function. This is only included if the top level exported object is a Keras model. """ class LayerAttributes(SerializedAttributes.with_attributes( 'LayerAttributes', checkpointable_objects=['non_trainable_variables', 'layers', 'metrics', 'layer_regularization_losses', 'layer_metrics'], functions=['call_and_return_conditional_losses', 'activity_regularizer_fn'], copy_from=[CommonEndpoints] )): """Layer checkpointable objects + functions that are saved to the SavedModel. List of all attributes: All attributes from CommonEndpoints non_trainable_variables: List of non-trainable variables in the layer and its sublayers. layers: List of all sublayers. metrics: List of all metrics in the layer and its sublayers. call_and_return_conditional_losses: Function that takes inputs and returns a tuple of (outputs of the call function, list of input-dependent losses). The list of losses excludes the activity regularizer function, which is separate to allow the deserialized Layer object to define a different activity regularizer. activity_regularizer_fn: Callable that returns the activity regularizer loss layer_regularization_losses: List of losses owned only by this layer. layer_metrics: List of metrics owned by this layer. """ class ModelAttributes(SerializedAttributes.with_attributes( 'ModelAttributes', copy_from=[LayerAttributes])): """Model checkpointable objects + functions that are saved to the SavedModel. List of all attributes: All attributes from LayerAttributes (including CommonEndpoints) """ # TODO(kathywu): Add attributes `compile_losses` and `compile_metrics`, which # list all losses and metrics defined by `model.compile`. class MetricAttributes( SerializedAttributes.with_attributes( 'MetricAttributes', checkpointable_objects=['variables'], functions=[], )): """Attributes that are added to Metric objects when saved to SavedModel. List of all attributes: variables: list of all variables """ pass class RNNAttributes(SerializedAttributes.with_attributes( 'RNNAttributes', checkpointable_objects=['states'], copy_from=[LayerAttributes])): """RNN checkpointable objects + functions that are saved to the SavedModel. List of all attributes: All attributes from LayerAttributes (including CommonEndpoints) states: List of state variables """