372 lines
13 KiB
Python
372 lines
13 KiB
Python
|
# Copyright 2015 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Utils related to keras model saving."""
|
||
|
|
|
||
|
|
import copy
|
||
|
|
import os
|
||
|
|
|
||
|
|
import tensorflow.compat.v2 as tf
|
||
|
|
|
||
|
|
import keras
|
||
|
|
from keras import backend
|
||
|
|
from keras import losses
|
||
|
|
from keras import optimizers
|
||
|
|
from keras.engine import base_layer_utils
|
||
|
|
from keras.optimizers import optimizer_v1
|
||
|
|
from keras.saving.legacy import serialization
|
||
|
|
from keras.utils import version_utils
|
||
|
|
from keras.utils.io_utils import ask_to_proceed_with_overwrite
|
||
|
|
|
||
|
|
# isort: off
|
||
|
|
from tensorflow.python.platform import tf_logging as logging
|
||
|
|
|
||
|
|
|
||
|
|
def extract_model_metrics(model):
|
||
|
|
"""Convert metrics from a Keras model `compile` API to dictionary.
|
||
|
|
|
||
|
|
This is used for converting Keras models to Estimators and SavedModels.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: A `tf.keras.Model` object.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Dictionary mapping metric names to metric instances. May return `None` if
|
||
|
|
the model does not contain any metrics.
|
||
|
|
"""
|
||
|
|
if getattr(model, "_compile_metrics", None):
|
||
|
|
# TODO(psv/kathywu): use this implementation in model to estimator flow.
|
||
|
|
# We are not using model.metrics here because we want to exclude the
|
||
|
|
# metrics added using `add_metric` API.
|
||
|
|
return {m.name: m for m in model._compile_metric_functions}
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def model_call_inputs(model, keep_original_batch_size=False):
|
||
|
|
"""Inspect model to get its input signature.
|
||
|
|
|
||
|
|
The model's input signature is a list with a single (possibly-nested)
|
||
|
|
object. This is due to the Keras-enforced restriction that tensor inputs
|
||
|
|
must be passed in as the first argument.
|
||
|
|
|
||
|
|
For example, a model with input {'feature1': <Tensor>, 'feature2': <Tensor>}
|
||
|
|
will have input signature:
|
||
|
|
[{'feature1': TensorSpec, 'feature2': TensorSpec}]
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: Keras Model object.
|
||
|
|
keep_original_batch_size: A boolean indicating whether we want to keep
|
||
|
|
using the original batch size or set it to None. Default is `False`,
|
||
|
|
which means that the batch dim of the returned input signature will
|
||
|
|
always be set to `None`.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tuple containing `(args, kwargs)` TensorSpecs of the model call function
|
||
|
|
inputs.
|
||
|
|
`kwargs` does not contain the `training` argument.
|
||
|
|
"""
|
||
|
|
input_specs = model.save_spec(dynamic_batch=not keep_original_batch_size)
|
||
|
|
if input_specs is None:
|
||
|
|
return None, None
|
||
|
|
input_specs = _enforce_names_consistency(input_specs)
|
||
|
|
return input_specs
|
||
|
|
|
||
|
|
|
||
|
|
def raise_model_input_error(model):
|
||
|
|
if isinstance(model, keras.models.Sequential):
|
||
|
|
raise ValueError(
|
||
|
|
f"Model {model} cannot be saved because the input shape is not "
|
||
|
|
"available. Please specify an input shape either by calling "
|
||
|
|
"`build(input_shape)` directly, or by calling the model on actual "
|
||
|
|
"data using `Model()`, `Model.fit()`, or `Model.predict()`."
|
||
|
|
)
|
||
|
|
|
||
|
|
# If the model is not a `Sequential`, it is intended to be a subclassed
|
||
|
|
# model.
|
||
|
|
raise ValueError(
|
||
|
|
f"Model {model} cannot be saved either because the input shape is not "
|
||
|
|
"available or because the forward pass of the model is not defined."
|
||
|
|
"To define a forward pass, please override `Model.call()`. To specify "
|
||
|
|
"an input shape, either call `build(input_shape)` directly, or call "
|
||
|
|
"the model on actual data using `Model()`, `Model.fit()`, or "
|
||
|
|
"`Model.predict()`. If you have a custom training step, please make "
|
||
|
|
"sure to invoke the forward pass in train step through "
|
||
|
|
"`Model.__call__`, i.e. `model(inputs)`, as opposed to `model.call()`."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def trace_model_call(model, input_signature=None):
|
||
|
|
"""Trace the model call to create a tf.function for exporting a Keras model.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: A Keras model.
|
||
|
|
input_signature: optional, a list of tf.TensorSpec objects specifying the
|
||
|
|
inputs to the model.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tf.function wrapping the model's call function with input signatures
|
||
|
|
set.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: if input signature cannot be inferred from the model.
|
||
|
|
"""
|
||
|
|
if input_signature is None:
|
||
|
|
if isinstance(model.call, tf.__internal__.function.Function):
|
||
|
|
input_signature = model.call.input_signature
|
||
|
|
|
||
|
|
if input_signature:
|
||
|
|
model_args = input_signature
|
||
|
|
model_kwargs = {}
|
||
|
|
else:
|
||
|
|
model_args, model_kwargs = model_call_inputs(model)
|
||
|
|
|
||
|
|
if model_args is None:
|
||
|
|
raise_model_input_error(model)
|
||
|
|
|
||
|
|
@tf.function
|
||
|
|
def _wrapped_model(*args, **kwargs):
|
||
|
|
"""A concrete tf.function that wraps the model's call function."""
|
||
|
|
(args, kwargs,) = model._call_spec.set_arg_value(
|
||
|
|
"training", False, args, kwargs, inputs_in_args=True
|
||
|
|
)
|
||
|
|
|
||
|
|
with base_layer_utils.call_context().enter(
|
||
|
|
model, inputs=None, build_graph=False, training=False, saving=True
|
||
|
|
):
|
||
|
|
outputs = model(*args, **kwargs)
|
||
|
|
|
||
|
|
# Outputs always has to be a flat dict.
|
||
|
|
output_names = model.output_names # Functional Model.
|
||
|
|
if output_names is None: # Subclassed Model.
|
||
|
|
from keras.engine import compile_utils
|
||
|
|
|
||
|
|
output_names = compile_utils.create_pseudo_output_names(outputs)
|
||
|
|
outputs = tf.nest.flatten(outputs)
|
||
|
|
return {name: output for name, output in zip(output_names, outputs)}
|
||
|
|
|
||
|
|
return _wrapped_model.get_concrete_function(*model_args, **model_kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def model_metadata(model, include_optimizer=True, require_config=True):
|
||
|
|
"""Returns a dictionary containing the model metadata."""
|
||
|
|
from keras import __version__ as keras_version
|
||
|
|
from keras.optimizers.legacy import optimizer_v2
|
||
|
|
|
||
|
|
model_config = {"class_name": model.__class__.__name__}
|
||
|
|
try:
|
||
|
|
model_config["config"] = model.get_config()
|
||
|
|
except NotImplementedError as e:
|
||
|
|
if require_config:
|
||
|
|
raise e
|
||
|
|
|
||
|
|
metadata = dict(
|
||
|
|
keras_version=str(keras_version),
|
||
|
|
backend=backend.backend(),
|
||
|
|
model_config=model_config,
|
||
|
|
)
|
||
|
|
if model.optimizer and include_optimizer:
|
||
|
|
if isinstance(model.optimizer, optimizer_v1.TFOptimizer):
|
||
|
|
logging.warning(
|
||
|
|
"TensorFlow optimizers do not "
|
||
|
|
"make it possible to access "
|
||
|
|
"optimizer attributes or optimizer state "
|
||
|
|
"after instantiation. "
|
||
|
|
"As a result, we cannot save the optimizer "
|
||
|
|
"as part of the model save file. "
|
||
|
|
"You will have to compile your model again after loading it. "
|
||
|
|
"Prefer using a Keras optimizer instead "
|
||
|
|
"(see keras.io/optimizers)."
|
||
|
|
)
|
||
|
|
elif model._compile_was_called:
|
||
|
|
training_config = model._get_compile_args(user_metrics=False)
|
||
|
|
training_config.pop("optimizer", None) # Handled separately.
|
||
|
|
metadata["training_config"] = _serialize_nested_config(
|
||
|
|
training_config
|
||
|
|
)
|
||
|
|
if isinstance(model.optimizer, optimizer_v2.RestoredOptimizer):
|
||
|
|
raise NotImplementedError(
|
||
|
|
"Optimizers loaded from a SavedModel cannot be saved. "
|
||
|
|
"If you are calling `model.save` or "
|
||
|
|
"`tf.keras.models.save_model`, "
|
||
|
|
"please set the `include_optimizer` option to `False`. For "
|
||
|
|
"`tf.saved_model.save`, "
|
||
|
|
"delete the optimizer from the model."
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
optimizer_config = {
|
||
|
|
"class_name": keras.utils.get_registered_name(
|
||
|
|
model.optimizer.__class__
|
||
|
|
),
|
||
|
|
"config": model.optimizer.get_config(),
|
||
|
|
}
|
||
|
|
metadata["training_config"]["optimizer_config"] = optimizer_config
|
||
|
|
return metadata
|
||
|
|
|
||
|
|
|
||
|
|
def should_overwrite(filepath, overwrite):
|
||
|
|
"""Returns whether the filepath should be overwritten."""
|
||
|
|
# If file exists and should not be overwritten.
|
||
|
|
if not overwrite and os.path.isfile(filepath):
|
||
|
|
return ask_to_proceed_with_overwrite(filepath)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def compile_args_from_training_config(training_config, custom_objects=None):
|
||
|
|
"""Return model.compile arguments from training config."""
|
||
|
|
if custom_objects is None:
|
||
|
|
custom_objects = {}
|
||
|
|
|
||
|
|
with keras.utils.CustomObjectScope(custom_objects):
|
||
|
|
optimizer_config = training_config["optimizer_config"]
|
||
|
|
optimizer = optimizers.deserialize(optimizer_config)
|
||
|
|
|
||
|
|
# Recover losses.
|
||
|
|
loss = None
|
||
|
|
loss_config = training_config.get("loss", None)
|
||
|
|
if loss_config is not None:
|
||
|
|
loss = _deserialize_nested_config(losses.deserialize, loss_config)
|
||
|
|
|
||
|
|
# Recover metrics.
|
||
|
|
metrics = None
|
||
|
|
metrics_config = training_config.get("metrics", None)
|
||
|
|
if metrics_config is not None:
|
||
|
|
metrics = _deserialize_nested_config(
|
||
|
|
_deserialize_metric, metrics_config
|
||
|
|
)
|
||
|
|
|
||
|
|
# Recover weighted metrics.
|
||
|
|
weighted_metrics = None
|
||
|
|
weighted_metrics_config = training_config.get("weighted_metrics", None)
|
||
|
|
if weighted_metrics_config is not None:
|
||
|
|
weighted_metrics = _deserialize_nested_config(
|
||
|
|
_deserialize_metric, weighted_metrics_config
|
||
|
|
)
|
||
|
|
|
||
|
|
sample_weight_mode = (
|
||
|
|
training_config["sample_weight_mode"]
|
||
|
|
if hasattr(training_config, "sample_weight_mode")
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
loss_weights = training_config["loss_weights"]
|
||
|
|
|
||
|
|
return dict(
|
||
|
|
optimizer=optimizer,
|
||
|
|
loss=loss,
|
||
|
|
metrics=metrics,
|
||
|
|
weighted_metrics=weighted_metrics,
|
||
|
|
loss_weights=loss_weights,
|
||
|
|
sample_weight_mode=sample_weight_mode,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _deserialize_nested_config(deserialize_fn, config):
|
||
|
|
"""Deserializes arbitrary Keras `config` using `deserialize_fn`."""
|
||
|
|
|
||
|
|
def _is_single_object(obj):
|
||
|
|
if isinstance(obj, dict) and "class_name" in obj:
|
||
|
|
return True # Serialized Keras object.
|
||
|
|
if isinstance(obj, str):
|
||
|
|
return True # Serialized function or string.
|
||
|
|
return False
|
||
|
|
|
||
|
|
if config is None:
|
||
|
|
return None
|
||
|
|
if _is_single_object(config):
|
||
|
|
return deserialize_fn(config)
|
||
|
|
elif isinstance(config, dict):
|
||
|
|
return {
|
||
|
|
k: _deserialize_nested_config(deserialize_fn, v)
|
||
|
|
for k, v in config.items()
|
||
|
|
}
|
||
|
|
elif isinstance(config, (tuple, list)):
|
||
|
|
return [
|
||
|
|
_deserialize_nested_config(deserialize_fn, obj) for obj in config
|
||
|
|
]
|
||
|
|
|
||
|
|
raise ValueError(
|
||
|
|
"Saved configuration not understood. Configuration should be a "
|
||
|
|
f"dictionary, string, tuple or list. Received: config={config}."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _serialize_nested_config(config):
|
||
|
|
"""Serialized a nested structure of Keras objects."""
|
||
|
|
|
||
|
|
def _serialize_fn(obj):
|
||
|
|
if callable(obj):
|
||
|
|
return serialization.serialize_keras_object(obj)
|
||
|
|
return obj
|
||
|
|
|
||
|
|
return tf.nest.map_structure(_serialize_fn, config)
|
||
|
|
|
||
|
|
|
||
|
|
def _deserialize_metric(metric_config):
|
||
|
|
"""Deserialize metrics, leaving special strings untouched."""
|
||
|
|
from keras import metrics as metrics_module
|
||
|
|
|
||
|
|
if metric_config in ["accuracy", "acc", "crossentropy", "ce"]:
|
||
|
|
# Do not deserialize accuracy and cross-entropy strings as we have
|
||
|
|
# special case handling for these in compile, based on model output
|
||
|
|
# shape.
|
||
|
|
return metric_config
|
||
|
|
return metrics_module.deserialize(metric_config)
|
||
|
|
|
||
|
|
|
||
|
|
def _enforce_names_consistency(specs):
|
||
|
|
"""Enforces that either all specs have names or none do."""
|
||
|
|
|
||
|
|
def _has_name(spec):
|
||
|
|
return spec is None or (hasattr(spec, "name") and spec.name is not None)
|
||
|
|
|
||
|
|
def _clear_name(spec):
|
||
|
|
spec = copy.deepcopy(spec)
|
||
|
|
if hasattr(spec, "name"):
|
||
|
|
spec._name = None
|
||
|
|
return spec
|
||
|
|
|
||
|
|
flat_specs = tf.nest.flatten(specs)
|
||
|
|
name_inconsistency = any(_has_name(s) for s in flat_specs) and not all(
|
||
|
|
_has_name(s) for s in flat_specs
|
||
|
|
)
|
||
|
|
|
||
|
|
if name_inconsistency:
|
||
|
|
specs = tf.nest.map_structure(_clear_name, specs)
|
||
|
|
return specs
|
||
|
|
|
||
|
|
|
||
|
|
def try_build_compiled_arguments(model):
|
||
|
|
if (
|
||
|
|
not version_utils.is_v1_layer_or_model(model)
|
||
|
|
and model.outputs is not None
|
||
|
|
):
|
||
|
|
try:
|
||
|
|
if not model.compiled_loss.built:
|
||
|
|
model.compiled_loss.build(model.outputs)
|
||
|
|
if not model.compiled_metrics.built:
|
||
|
|
model.compiled_metrics.build(model.outputs, model.outputs)
|
||
|
|
except: # noqa: E722
|
||
|
|
logging.warning(
|
||
|
|
"Compiled the loaded model, but the compiled metrics have "
|
||
|
|
"yet to be built. `model.compile_metrics` will be empty "
|
||
|
|
"until you train or evaluate the model."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def is_hdf5_filepath(filepath):
|
||
|
|
return (
|
||
|
|
filepath.endswith(".h5")
|
||
|
|
or filepath.endswith(".keras")
|
||
|
|
or filepath.endswith(".hdf5")
|
||
|
|
)
|