891 lines
35 KiB
Python
891 lines
35 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.
|
|
# ==============================================================================
|
|
|
|
"""Code for model cloning, plus model-related API entries."""
|
|
|
|
import tensorflow.compat.v2 as tf
|
|
|
|
from keras import backend
|
|
from keras import metrics as metrics_module
|
|
from keras.engine import functional
|
|
from keras.engine import sequential
|
|
from keras.engine import training
|
|
from keras.engine import training_v1
|
|
from keras.engine.base_layer import AddMetric
|
|
from keras.engine.base_layer import Layer
|
|
from keras.engine.input_layer import Input
|
|
from keras.engine.input_layer import InputLayer
|
|
from keras.optimizers import optimizer_v1
|
|
from keras.saving.legacy import serialization
|
|
from keras.saving.object_registration import CustomObjectScope
|
|
from keras.utils import generic_utils
|
|
from keras.utils import version_utils
|
|
|
|
# isort: off
|
|
from tensorflow.python.platform import tf_logging as logging
|
|
from tensorflow.python.util.tf_export import keras_export
|
|
|
|
# API entries importable from `keras.models`:
|
|
Model = training.Model
|
|
Sequential = sequential.Sequential
|
|
|
|
|
|
# Callable used to clone a layer with weights preserved.
|
|
def share_weights(layer):
|
|
return layer
|
|
|
|
|
|
def _clone_layer(layer):
|
|
return layer.__class__.from_config(layer.get_config())
|
|
|
|
|
|
def _insert_ancillary_layers(model, ancillary_layers, metrics_names, new_nodes):
|
|
"""Inserts ancillary layers into the model with the proper order."""
|
|
# Sort `AddMetric` layers so they agree with metrics_names.
|
|
metric_layers = [
|
|
layer for layer in ancillary_layers if isinstance(layer, AddMetric)
|
|
]
|
|
metric_layers.sort(key=lambda layer: metrics_names.index(layer.metric_name))
|
|
ancillary_layers = [
|
|
layer for layer in ancillary_layers if not isinstance(layer, AddMetric)
|
|
] + metric_layers
|
|
model._insert_layers(ancillary_layers, relevant_nodes=list(new_nodes))
|
|
|
|
|
|
def _make_new_nodes(nodes_by_depth, layer_fn, layer_map, tensor_map):
|
|
"""Make new nodes with the layers in `layer_map` based on `nodes_by_depth`.
|
|
|
|
Args:
|
|
nodes_by_depth: Provides structure information to create new nodes.
|
|
layer_fn: Function to clone layers.
|
|
layer_map: Map from layers in `model` to new layers.
|
|
tensor_map: Map from tensors in `model` to newly compute tensors.
|
|
|
|
Returns:
|
|
A set of new nodes. `layer_map` and `tensor_map` are updated.
|
|
"""
|
|
# Iterated over every node in the reference model, in depth order.
|
|
new_nodes = set()
|
|
depth_keys = list(nodes_by_depth.keys())
|
|
depth_keys.sort(reverse=True)
|
|
for depth in depth_keys:
|
|
nodes = nodes_by_depth[depth]
|
|
for node in nodes:
|
|
# Recover the corresponding layer.
|
|
layer = node.outbound_layer
|
|
|
|
# Get or create layer.
|
|
if layer not in layer_map:
|
|
new_layer = layer_fn(layer)
|
|
layer_map[layer] = new_layer
|
|
layer = new_layer
|
|
else:
|
|
# Reuse previously cloned layer.
|
|
layer = layer_map[layer]
|
|
# Don't call InputLayer multiple times.
|
|
if isinstance(layer, InputLayer):
|
|
continue
|
|
|
|
# If all previous input tensors are available in tensor_map,
|
|
# then call node.inbound_layer on them.
|
|
if all(
|
|
tensor in tensor_map
|
|
for tensor in tf.nest.flatten(node.input_tensors)
|
|
):
|
|
# Call layer.
|
|
args = tf.nest.map_structure(
|
|
lambda t: tensor_map.get(t, t), node.call_args
|
|
)
|
|
kwargs = tf.nest.map_structure(
|
|
lambda t: tensor_map.get(t, t), node.call_kwargs
|
|
)
|
|
output_tensors = layer(*args, **kwargs)
|
|
|
|
# Thread-safe way to keep track of what node was created.
|
|
first_output_tensor = tf.nest.flatten(output_tensors)[0]
|
|
new_nodes.add(
|
|
layer._inbound_nodes[
|
|
first_output_tensor._keras_history.node_index
|
|
]
|
|
)
|
|
|
|
for x, y in zip(
|
|
tf.nest.flatten(node.output_tensors),
|
|
tf.nest.flatten(output_tensors),
|
|
):
|
|
tensor_map[x] = y
|
|
return new_nodes
|
|
|
|
|
|
def _clone_functional_model(model, input_tensors=None, layer_fn=_clone_layer):
|
|
"""Clone a functional `Model` instance.
|
|
|
|
Model cloning is similar to calling a model on new inputs,
|
|
except that it creates new layers (and thus new weights) instead
|
|
of sharing the weights of the existing layers.
|
|
|
|
Input layers are always cloned.
|
|
|
|
Args:
|
|
model: Instance of `Model`.
|
|
input_tensors: optional list of input tensors
|
|
to build the model upon. If not provided,
|
|
placeholders will be created.
|
|
layer_fn: callable to be applied on non-input layers in the model. By
|
|
default it clones the layer. Another example is to preserve the
|
|
layer to share the weights. This is required when we create a
|
|
per-replica copy of the model with distribution strategy; we want
|
|
the weights to be shared but still feed inputs separately so we
|
|
create new input layers.
|
|
|
|
Returns:
|
|
An instance of `Model` reproducing the behavior
|
|
of the original model, on top of new inputs tensors,
|
|
using newly instantiated weights.
|
|
|
|
Raises:
|
|
ValueError: in case of invalid `model` argument value or `layer_fn`
|
|
argument value.
|
|
"""
|
|
if layer_fn is None:
|
|
layer_fn = _clone_layer
|
|
|
|
if not isinstance(model, Model):
|
|
raise ValueError(
|
|
"Expected `model` argument "
|
|
f"to be a `Model` instance. Received: model={model}"
|
|
)
|
|
if isinstance(model, Sequential):
|
|
raise ValueError(
|
|
"Expected `model` argument "
|
|
"to be a functional `Model` instance, "
|
|
f"got a `Sequential` instance instead: {model}"
|
|
)
|
|
if not model._is_graph_network:
|
|
raise ValueError(
|
|
"Expected `model` argument "
|
|
"to be a functional `Model` instance, "
|
|
f"but got a subclassed model instead: {model}"
|
|
)
|
|
|
|
new_input_layers = {} # Cache for created layers.
|
|
if input_tensors is not None:
|
|
# Make sure that all input tensors come from a Keras layer.
|
|
input_tensors = tf.nest.flatten(input_tensors)
|
|
for i, input_tensor in enumerate(input_tensors):
|
|
original_input_layer = model._input_layers[i]
|
|
|
|
# Cache input layer. Create a new layer if the tensor is originally
|
|
# not from a Keras layer.
|
|
if not backend.is_keras_tensor(input_tensor):
|
|
name = original_input_layer.name
|
|
input_tensor = Input(
|
|
tensor=input_tensor, name="input_wrapper_for_" + name
|
|
)
|
|
newly_created_input_layer = input_tensor._keras_history.layer
|
|
new_input_layers[
|
|
original_input_layer
|
|
] = newly_created_input_layer
|
|
else:
|
|
new_input_layers[
|
|
original_input_layer
|
|
] = input_tensor._keras_history.layer
|
|
|
|
if not callable(layer_fn):
|
|
raise ValueError(
|
|
"Expected `layer_fn` argument to be a callable. "
|
|
f"Received: layer_fn={layer_fn}"
|
|
)
|
|
|
|
model_configs, created_layers = _clone_layers_and_model_config(
|
|
model, new_input_layers, layer_fn
|
|
)
|
|
# Reconstruct model from the config, using the cloned layers.
|
|
(
|
|
input_tensors,
|
|
output_tensors,
|
|
created_layers,
|
|
) = functional.reconstruct_from_config(
|
|
model_configs, created_layers=created_layers
|
|
)
|
|
metrics_names = model.metrics_names
|
|
if functional.has_functional_like_constructor(model.__class__):
|
|
new_model = model.__class__(
|
|
input_tensors, output_tensors, name=model.name
|
|
)
|
|
else:
|
|
# This may be incorrect: the new model will end up having a different
|
|
# class than the original. However various existing models rely
|
|
# on this behavior, so we keep it.
|
|
new_model = Model(input_tensors, output_tensors, name=model.name)
|
|
|
|
# Layers not directly tied to outputs of the Model, such as loss layers
|
|
# created in `add_loss` and `add_metric`.
|
|
ancillary_layers = [
|
|
layer
|
|
for layer in created_layers.values()
|
|
if layer not in new_model.layers
|
|
]
|
|
# TODO(b/162887610): This may need to adjust the inbound node index if the
|
|
# created layers had already been used to define other models.
|
|
if ancillary_layers:
|
|
new_nodes = tf.nest.flatten(
|
|
[
|
|
layer.inbound_nodes[1:]
|
|
if functional._should_skip_first_node(layer)
|
|
else layer.inbound_nodes
|
|
for layer in created_layers.values()
|
|
]
|
|
)
|
|
_insert_ancillary_layers(
|
|
new_model, ancillary_layers, metrics_names, new_nodes
|
|
)
|
|
return new_model
|
|
|
|
|
|
def _clone_layers_and_model_config(model, input_layers, layer_fn):
|
|
"""Clones all layers; returns the model config without serializing layers.
|
|
|
|
This function ensures that only the node graph is retrieved when getting the
|
|
model config. The `layer_fn` used to clone layers might not rely on
|
|
`layer.get_config()`, so some custom layers do not define `get_config`.
|
|
Trying to retrieve the config results in errors.
|
|
|
|
Args:
|
|
model: A Functional model.
|
|
input_layers: Dictionary mapping input layers in `model` to new input
|
|
layers.
|
|
layer_fn: Function used to clone all non-input layers.
|
|
|
|
Returns:
|
|
Model config object, and a dictionary of newly created layers.
|
|
"""
|
|
created_layers = {}
|
|
|
|
def _copy_layer(layer):
|
|
# Whenever the network config attempts to get the layer serialization,
|
|
# return a dummy dictionary.
|
|
if layer in input_layers:
|
|
created_layers[layer.name] = input_layers[layer]
|
|
elif layer in model._input_layers:
|
|
created_layers[layer.name] = InputLayer(**layer.get_config())
|
|
else:
|
|
created_layers[layer.name] = layer_fn(layer)
|
|
return {}
|
|
|
|
config = functional.get_network_config(
|
|
model, serialize_layer_fn=_copy_layer
|
|
)
|
|
return config, created_layers
|
|
|
|
|
|
def _remove_ancillary_layers(model, layer_map, layers):
|
|
"""Removes and returns any ancillary layers from `layers` based on `model`.
|
|
|
|
Ancillary layers are part of the model topology but not used to compute the
|
|
model outputs, e.g., layers from `add_loss` and `add_metric`.
|
|
|
|
Args:
|
|
model: A Keras Model.
|
|
layer_map: A map to from layers in the `model` to those in `layers`.
|
|
layers: A list of all layers.
|
|
|
|
Returns:
|
|
Two lists of layers: (1) `layers` with the ancillary layers removed, and
|
|
(2) the ancillary layers.
|
|
"""
|
|
ancillary_layers = [] # Additional layers for computing losses and metrics.
|
|
if not model._is_graph_network:
|
|
return layers, ancillary_layers
|
|
|
|
# Ancillary layers are those with depth < 0.
|
|
depths = [depth for depth in model._nodes_by_depth.keys() if depth < 0]
|
|
depths.sort(reverse=True) # Order topologically from inputs to outputs.
|
|
for depth in depths:
|
|
for node in model._nodes_by_depth[depth]:
|
|
ancillary_layers.append(layer_map[node.outbound_layer])
|
|
|
|
return [l for l in layers if l not in ancillary_layers], ancillary_layers
|
|
|
|
|
|
def _clone_sequential_model(model, input_tensors=None, layer_fn=_clone_layer):
|
|
"""Clone a `Sequential` model instance.
|
|
|
|
Model cloning is similar to calling a model on new inputs,
|
|
except that it creates new layers (and thus new weights) instead
|
|
of sharing the weights of the existing layers.
|
|
|
|
Args:
|
|
model: Instance of `Sequential`.
|
|
input_tensors: optional list of input tensors
|
|
to build the model upon. If not provided,
|
|
placeholders will be created.
|
|
layer_fn: callable to be applied on non-input layers in the model. By
|
|
default it clones the layer. Another example is to preserve the
|
|
layer to share the weights. This is required when we create a
|
|
per-replica copy of the model with distribution strategy; we want
|
|
the weights to be shared but still feed inputs separately so we
|
|
create new input layers.
|
|
|
|
Returns:
|
|
An instance of `Sequential` reproducing the behavior
|
|
of the original model, on top of new inputs tensors,
|
|
using newly instantiated weights.
|
|
|
|
Raises:
|
|
ValueError: in case of invalid `model` argument value or `layer_fn`
|
|
argument value.
|
|
"""
|
|
if layer_fn is None:
|
|
layer_fn = _clone_layer
|
|
|
|
if not isinstance(model, Sequential):
|
|
raise ValueError(
|
|
"Expected `model` argument "
|
|
"to be a `Sequential` model instance. "
|
|
f"Received: model={model}"
|
|
)
|
|
|
|
if not callable(layer_fn):
|
|
raise ValueError(
|
|
"Expected `layer_fn` argument to be a callable. "
|
|
f"Received: layer_fn={layer_fn}"
|
|
)
|
|
|
|
layers = [] # Layers needed to compute the model's outputs.
|
|
layer_map = {}
|
|
# Ensure that all layers are cloned. The model's layers
|
|
# property will exclude the initial InputLayer (if it exists) in the model,
|
|
# resulting in a different Sequential model structure.
|
|
for layer in model._flatten_layers(include_self=False, recursive=False):
|
|
if isinstance(layer, InputLayer) and input_tensors is not None:
|
|
# If input tensors are provided, the original model's InputLayer is
|
|
# overwritten with a different InputLayer.
|
|
continue
|
|
cloned_layer = (
|
|
_clone_layer(layer)
|
|
if isinstance(layer, InputLayer)
|
|
else layer_fn(layer)
|
|
)
|
|
layers.append(cloned_layer)
|
|
layer_map[layer] = cloned_layer
|
|
layers, ancillary_layers = _remove_ancillary_layers(
|
|
model, layer_map, layers
|
|
)
|
|
|
|
if input_tensors is None:
|
|
cloned_model = Sequential(layers=layers, name=model.name)
|
|
elif len(generic_utils.to_list(input_tensors)) != 1:
|
|
raise ValueError(
|
|
"To clone a `Sequential` model, we expect at most one tensor as "
|
|
f"part of `input_tensors`. Received: input_tensors={input_tensors}"
|
|
)
|
|
else:
|
|
# Overwrite the original model's input layer.
|
|
if isinstance(input_tensors, tuple):
|
|
input_tensors = list(input_tensors)
|
|
x = generic_utils.to_list(input_tensors)[0]
|
|
if backend.is_keras_tensor(x):
|
|
origin_layer = x._keras_history.layer
|
|
if isinstance(origin_layer, InputLayer):
|
|
cloned_model = Sequential(
|
|
layers=[origin_layer] + layers, name=model.name
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"Cannot clone a `Sequential` model on top "
|
|
"of a tensor that comes from a Keras layer "
|
|
"other than an `InputLayer`. "
|
|
"Use the Functional API instead. "
|
|
f"Received: input_tensors={input_tensors}"
|
|
)
|
|
else:
|
|
input_tensor = Input(
|
|
tensor=x, name="input_wrapper_for_" + str(x.name)
|
|
)
|
|
input_layer = input_tensor._keras_history.layer
|
|
cloned_model = Sequential(
|
|
layers=[input_layer] + layers, name=model.name
|
|
)
|
|
|
|
if not ancillary_layers:
|
|
return cloned_model
|
|
|
|
tensor_map = {} # Maps tensors from `model` to those in `cloned_model`.
|
|
for depth, cloned_nodes in cloned_model._nodes_by_depth.items():
|
|
nodes = model._nodes_by_depth[depth]
|
|
# This should be safe in a Sequential model. In an arbitrary network,
|
|
# you need to sort using the outbound layer of the node as a key.
|
|
for cloned_node, node in zip(cloned_nodes, nodes):
|
|
if isinstance(cloned_node.output_tensors, list):
|
|
for j, output_tensor in enumerate(cloned_node.output_tensors):
|
|
tensor_map[node.output_tensors[j]] = output_tensor
|
|
else:
|
|
tensor_map[node.output_tensors] = cloned_node.output_tensors
|
|
# Ancillary nodes have negative depth.
|
|
new_nodes = _make_new_nodes(
|
|
{
|
|
depth: nodes
|
|
for depth, nodes in model._nodes_by_depth.items()
|
|
if depth < 0
|
|
},
|
|
layer_fn,
|
|
layer_map,
|
|
tensor_map,
|
|
)
|
|
_insert_ancillary_layers(
|
|
cloned_model, ancillary_layers, model.metrics_names, new_nodes
|
|
)
|
|
return cloned_model
|
|
|
|
|
|
@keras_export("keras.models.clone_model")
|
|
def clone_model(model, input_tensors=None, clone_function=None):
|
|
"""Clone a Functional or Sequential `Model` instance.
|
|
|
|
Model cloning is similar to calling a model on new inputs,
|
|
except that it creates new layers (and thus new weights) instead
|
|
of sharing the weights of the existing layers.
|
|
|
|
Note that
|
|
`clone_model` will not preserve the uniqueness of shared objects within the
|
|
model (e.g. a single variable attached to two distinct layers will be
|
|
restored as two separate variables).
|
|
|
|
Args:
|
|
model: Instance of `Model`
|
|
(could be a Functional model or a Sequential model).
|
|
input_tensors: optional list of input tensors or InputLayer objects
|
|
to build the model upon. If not provided,
|
|
new `Input` objects will be created.
|
|
clone_function: Callable to be used to clone each layer in the target
|
|
model (except `InputLayer` instances). It takes as argument the
|
|
layer instance to be cloned, and returns the corresponding layer
|
|
instance to be used in the model copy. If unspecified, this callable
|
|
defaults to the following serialization/deserialization function:
|
|
`lambda layer: layer.__class__.from_config(layer.get_config())`.
|
|
By passing a custom callable, you can customize your copy of the
|
|
model, e.g. by wrapping certain layers of interest (you might want
|
|
to replace all `LSTM` instances with equivalent
|
|
`Bidirectional(LSTM(...))` instances, for example).
|
|
|
|
Returns:
|
|
An instance of `Model` reproducing the behavior
|
|
of the original model, on top of new inputs tensors,
|
|
using newly instantiated weights. The cloned model may behave
|
|
differently from the original model if a custom `clone_function`
|
|
modifies the layer.
|
|
|
|
Example:
|
|
|
|
```python
|
|
# Create a test Sequential model.
|
|
model = keras.Sequential([
|
|
keras.Input(shape=(728,)),
|
|
keras.layers.Dense(32, activation='relu'),
|
|
keras.layers.Dense(1, activation='sigmoid'),
|
|
])
|
|
# Create a copy of the test model (with freshly initialized weights).
|
|
new_model = clone_model(model)
|
|
```
|
|
|
|
Note that subclassed models cannot be cloned, since their internal
|
|
layer structure is not known. To achieve equivalent functionality
|
|
as `clone_model` in the case of a subclassed model, simply make sure
|
|
that the model class implements `get_config()`
|
|
(and optionally `from_config()`), and call:
|
|
|
|
```python
|
|
new_model = model.__class__.from_config(model.get_config())
|
|
```
|
|
"""
|
|
with serialization.DisableSharedObjectScope():
|
|
if isinstance(model, Sequential):
|
|
return _clone_sequential_model(
|
|
model, input_tensors=input_tensors, layer_fn=clone_function
|
|
)
|
|
if isinstance(model, functional.Functional):
|
|
# If the get_config() method is the same as a regular Functional
|
|
# model, we're safe to use _clone_functional_model (which relies
|
|
# on a Functional constructor). In the case where the get_config
|
|
# is custom, this may not necessarily work, but if clone_function
|
|
# or input_tensors are passed, we attempt it anyway
|
|
# in order to preserve backwards compatibility.
|
|
if generic_utils.is_default(model.get_config) or (
|
|
clone_function or input_tensors
|
|
):
|
|
return _clone_functional_model(
|
|
model, input_tensors=input_tensors, layer_fn=clone_function
|
|
)
|
|
|
|
# Case of a custom model class
|
|
if clone_function or input_tensors:
|
|
raise ValueError(
|
|
"Arguments clone_function and input_tensors "
|
|
"are only supported for Sequential models "
|
|
"or Functional models. Received model of "
|
|
f"type '{model.__class__.__name__}', with "
|
|
f"clone_function={clone_function} and "
|
|
f"input_tensors={input_tensors}"
|
|
)
|
|
# Note that a custom object scope may be required in this case.
|
|
return model.__class__.from_config(model.get_config())
|
|
|
|
|
|
# "Clone" a subclassed model by resetting all of the attributes.
|
|
def _in_place_subclassed_model_reset(model):
|
|
"""Substitute for model cloning that works for subclassed models.
|
|
|
|
Subclassed models cannot be cloned because their topology is not
|
|
serializable. To "instantiate" an identical model in a new TF graph, we
|
|
reuse the original model object, but we clear its state.
|
|
|
|
After calling this function on a model instance, you can use the model
|
|
instance as if it were a model clone (in particular you can use it in a new
|
|
graph).
|
|
|
|
This method clears the state of the input model. It is thus destructive.
|
|
However the original state can be restored fully by calling
|
|
`_in_place_subclassed_model_state_restoration`.
|
|
|
|
Args:
|
|
model: Instance of a Keras model created via subclassing.
|
|
|
|
Raises:
|
|
ValueError: In case the model uses a subclassed model as inner layer.
|
|
"""
|
|
assert (
|
|
not model._is_graph_network
|
|
) # Only makes sense for subclassed networks
|
|
# Select correct base class for new Model.
|
|
version_utils.swap_class(
|
|
model.__class__,
|
|
training.Model,
|
|
training_v1.Model,
|
|
tf.compat.v1.executing_eagerly_outside_functions(),
|
|
)
|
|
# Retrieve all layers tracked by the model as well as their attribute names
|
|
attributes_cache = {}
|
|
for name in dir(model):
|
|
# Skip attrs that track other trackables.
|
|
if name == "submodules" or name == "_self_tracked_trackables":
|
|
continue
|
|
|
|
try:
|
|
value = getattr(model, name)
|
|
except (AttributeError, ValueError, TypeError):
|
|
continue
|
|
if isinstance(value, Layer):
|
|
attributes_cache[name] = value
|
|
assert value in model.layers
|
|
if hasattr(value, "layers") and value.layers:
|
|
raise ValueError(
|
|
"We do not support the use of nested layers "
|
|
"in `model_to_estimator` at this time. Found nested "
|
|
f"layer: {value}"
|
|
)
|
|
elif isinstance(value, (list, tuple)) and name not in (
|
|
"layers",
|
|
"_layers",
|
|
"metrics",
|
|
"_compile_metric_functions",
|
|
"_output_loss_metrics",
|
|
):
|
|
# Handle case: list/tuple of layers (also tracked by the Network
|
|
# API).
|
|
if value and all(isinstance(val, Layer) for val in value):
|
|
raise ValueError(
|
|
"We do not support the use of list-of-layers "
|
|
"attributes in subclassed models used with "
|
|
"`model_to_estimator` at this time. Found list "
|
|
f"model: {name}"
|
|
)
|
|
|
|
# Replace layers on the model with fresh layers
|
|
layers_to_names = {value: key for key, value in attributes_cache.items()}
|
|
original_layers = list(
|
|
model._flatten_layers(include_self=False, recursive=False)
|
|
)
|
|
setattr_tracking = model._setattr_tracking
|
|
model._setattr_tracking = False
|
|
model._self_tracked_trackables = []
|
|
for layer in original_layers: # We preserve layer order.
|
|
config = layer.get_config()
|
|
# This will not work for nested subclassed models used as layers.
|
|
# This would be theoretically possible to support, but would add
|
|
# complexity. Only do it if users complain.
|
|
if isinstance(layer, training.Model) and not layer._is_graph_network:
|
|
raise ValueError(
|
|
"We do not support the use of nested subclassed models "
|
|
"in `model_to_estimator` at this time. Found nested "
|
|
f"model: {layer}"
|
|
)
|
|
fresh_layer = layer.__class__.from_config(config)
|
|
name = layers_to_names[layer]
|
|
setattr(model, name, fresh_layer)
|
|
model._self_tracked_trackables.append(fresh_layer)
|
|
|
|
# Cache original model build attributes (in addition to layers)
|
|
if (
|
|
not hasattr(model, "_original_attributes_cache")
|
|
or model._original_attributes_cache is None
|
|
):
|
|
if model.built:
|
|
attributes_to_cache = [
|
|
"inputs",
|
|
"outputs",
|
|
"total_loss",
|
|
"optimizer",
|
|
"train_function",
|
|
"test_function",
|
|
"predict_function",
|
|
"_training_endpoints",
|
|
"_collected_trainable_weights",
|
|
"_feed_inputs",
|
|
"_feed_input_names",
|
|
"_feed_input_shapes",
|
|
]
|
|
for name in attributes_to_cache:
|
|
attributes_cache[name] = getattr(model, name)
|
|
model._original_attributes_cache = attributes_cache
|
|
_reset_build_compile_trackers(model)
|
|
model._setattr_tracking = setattr_tracking
|
|
|
|
|
|
def _reset_build_compile_trackers(model):
|
|
"""Reset state trackers for model.
|
|
|
|
Note that we do not actually zero out attributes such as optimizer,
|
|
but instead rely on the expectation that all of the attrs will be
|
|
over-written on calling build/compile/etc. This is somewhat fragile,
|
|
insofar as we check elsewhere for the presence of these attributes as
|
|
evidence of having been built/compiled/etc. Pending a better way to do this,
|
|
we reset key attributes here to allow building and compiling.
|
|
|
|
Args:
|
|
model: the model that is being reset
|
|
"""
|
|
# Reset build state
|
|
model.built = False
|
|
model.inputs = None
|
|
model.outputs = None
|
|
# Reset compile state
|
|
model._is_compiled = False
|
|
if not tf.compat.v1.executing_eagerly_outside_functions():
|
|
model._v1_compile_was_called = False
|
|
model.optimizer = None
|
|
|
|
|
|
@keras_export(
|
|
"keras.__internal__.models.in_place_subclassed_model_state_restoration",
|
|
v1=[],
|
|
)
|
|
def in_place_subclassed_model_state_restoration(model):
|
|
"""Restores the original state of a model after it was "reset".
|
|
|
|
This undoes this action of `_in_place_subclassed_model_reset`, which is
|
|
called in `clone_and_build_model` if `in_place_reset` is set to True.
|
|
|
|
Args:
|
|
model: Instance of a Keras model created via subclassing, on which
|
|
`_in_place_subclassed_model_reset` was previously called.
|
|
"""
|
|
assert not model._is_graph_network
|
|
# Restore layers and build attributes
|
|
if (
|
|
hasattr(model, "_original_attributes_cache")
|
|
and model._original_attributes_cache is not None
|
|
):
|
|
# Models have sticky attribute assignment, so we want to be careful to
|
|
# add back the previous attributes and track Layers by their original
|
|
# names without adding dependencies on "utility" attributes which Models
|
|
# exempt when they're constructed.
|
|
setattr_tracking = model._setattr_tracking
|
|
model._setattr_tracking = False
|
|
model._self_tracked_trackables = []
|
|
for name, value in model._original_attributes_cache.items():
|
|
setattr(model, name, value)
|
|
if isinstance(value, Layer):
|
|
model._self_tracked_trackables.append(value)
|
|
model._original_attributes_cache = None
|
|
model._setattr_tracking = setattr_tracking
|
|
else:
|
|
# Restore to the state of a never-called model.
|
|
_reset_build_compile_trackers(model)
|
|
|
|
|
|
@keras_export("keras.__internal__.models.clone_and_build_model", v1=[])
|
|
def clone_and_build_model(
|
|
model,
|
|
input_tensors=None,
|
|
target_tensors=None,
|
|
custom_objects=None,
|
|
compile_clone=True,
|
|
in_place_reset=False,
|
|
optimizer_iterations=None,
|
|
optimizer_config=None,
|
|
):
|
|
"""Clone a `Model` and build/compile it with the same settings used before.
|
|
|
|
This function can be run in the same graph or in a separate graph from the
|
|
model. When using a separate graph, `in_place_reset` must be `False`.
|
|
|
|
Note that, currently, the clone produced from this function may not work
|
|
with TPU DistributionStrategy. Try at your own risk.
|
|
|
|
Args:
|
|
model: `tf.keras.Model` object. Can be Functional, Sequential, or
|
|
sub-classed.
|
|
input_tensors: Optional list or dictionary of input tensors to build the
|
|
model upon. If not provided, placeholders will be created.
|
|
target_tensors: Optional list of target tensors for compiling the model.
|
|
If not provided, placeholders will be created.
|
|
custom_objects: Optional dictionary mapping string names to custom classes
|
|
or functions.
|
|
compile_clone: Boolean, whether to compile model clone (default `True`).
|
|
in_place_reset: Boolean, whether to reset the model in place. Only used if
|
|
the model is a subclassed model. In the case of a subclassed model,
|
|
this argument must be set to `True` (default `False`). To restore the
|
|
original model, use the function
|
|
`in_place_subclassed_model_state_restoration(model)`.
|
|
optimizer_iterations: An iterations variable that will be incremented by
|
|
the optimizer if the clone is compiled. This argument is used when a
|
|
Keras model is cloned into an Estimator model function, because
|
|
Estimators create their own global step variable.
|
|
optimizer_config: Optimizer config dictionary or list of dictionary
|
|
returned from `get_config()`. This argument should be defined if
|
|
`clone_and_build_model` is called in a different graph or session from
|
|
the original model, and the optimizer is an instance of `OptimizerV2`.
|
|
|
|
Returns:
|
|
Clone of the model.
|
|
|
|
Raises:
|
|
ValueError: Cloning fails in the following cases
|
|
- cloning a subclassed model with `in_place_reset` set to False.
|
|
- compiling the clone when the original model has not been compiled.
|
|
"""
|
|
# Grab optimizer now, as we reset-in-place for subclassed models, but
|
|
# want to maintain access to the original optimizer.
|
|
orig_optimizer = model.optimizer
|
|
if compile_clone and not orig_optimizer:
|
|
raise ValueError(
|
|
"Error when cloning model: `compile_clone` was set to True, but "
|
|
f"the original model has not been compiled. Received: model={model}"
|
|
)
|
|
|
|
if compile_clone:
|
|
compile_args = model._get_compile_args()
|
|
# Allows this method to be robust to switching graph and eager classes.
|
|
model._get_compile_args = lambda: compile_args
|
|
|
|
with CustomObjectScope(custom_objects or {}):
|
|
if model._is_graph_network:
|
|
clone = clone_model(model, input_tensors=input_tensors)
|
|
elif isinstance(model, Sequential):
|
|
clone = clone_model(model, input_tensors=input_tensors)
|
|
if (
|
|
not clone._is_graph_network
|
|
and model._build_input_shape is not None
|
|
):
|
|
if tf.compat.v1.executing_eagerly_outside_functions():
|
|
clone.build(model._build_input_shape)
|
|
else:
|
|
clone._set_inputs(
|
|
backend.placeholder(
|
|
model._build_input_shape,
|
|
dtype=model.inputs[0].dtype,
|
|
)
|
|
)
|
|
else:
|
|
try:
|
|
# Prefer cloning the model if serial/deserial logic is
|
|
# implemented for subclassed model.
|
|
clone = model.__class__.from_config(model.get_config())
|
|
except NotImplementedError:
|
|
logging.warning(
|
|
"This model is a subclassed model. Please implement "
|
|
"`get_config` and `from_config` to better support "
|
|
"cloning the model."
|
|
)
|
|
if not in_place_reset:
|
|
raise ValueError(
|
|
f"This model ({model}) is a subclassed model. "
|
|
"Such a model cannot be cloned, but there is a "
|
|
"workaround where the model is reset in-place. "
|
|
"To use this, please set the "
|
|
"argument `in_place_reset` to `True`. This will reset "
|
|
"the attributes in the original model. "
|
|
"To restore the attributes, call "
|
|
"`in_place_subclassed_model_state_restoration(model)`."
|
|
)
|
|
clone = model
|
|
_in_place_subclassed_model_reset(clone)
|
|
if input_tensors is not None:
|
|
if (
|
|
isinstance(input_tensors, (list, tuple))
|
|
and len(input_tensors) == 1
|
|
):
|
|
input_tensors = input_tensors[0]
|
|
clone._set_inputs(input_tensors)
|
|
|
|
if compile_clone:
|
|
if isinstance(orig_optimizer, optimizer_v1.TFOptimizer):
|
|
optimizer = optimizer_v1.TFOptimizer(
|
|
orig_optimizer.optimizer, optimizer_iterations
|
|
)
|
|
backend.track_tf_optimizer(optimizer)
|
|
else:
|
|
if not isinstance(orig_optimizer, (tuple, list)):
|
|
orig_optimizer = [orig_optimizer]
|
|
if optimizer_config is None:
|
|
optimizer = [
|
|
opt.__class__.from_config(opt.get_config())
|
|
for opt in orig_optimizer
|
|
]
|
|
elif isinstance(optimizer_config, dict):
|
|
optimizer = [
|
|
orig_optimizer[0].__class__.from_config(optimizer_config)
|
|
]
|
|
else:
|
|
# optimizer config is list of dict, same order as
|
|
# orig_optimizer.
|
|
optimizer = [
|
|
opt.__class__.from_config(opt_config)
|
|
for (opt, opt_config) in zip(
|
|
orig_optimizer, optimizer_config
|
|
)
|
|
]
|
|
if optimizer_iterations is not None:
|
|
for opt in optimizer:
|
|
opt.iterations = optimizer_iterations
|
|
|
|
if len(optimizer) == 1:
|
|
optimizer = optimizer[0]
|
|
|
|
compile_args["optimizer"] = optimizer
|
|
if target_tensors is not None:
|
|
compile_args["target_tensors"] = target_tensors
|
|
# Ensure Metric objects in new model are separate from existing model.
|
|
compile_args["metrics"] = metrics_module.clone_metrics(
|
|
compile_args["metrics"]
|
|
)
|
|
compile_args["weighted_metrics"] = metrics_module.clone_metrics(
|
|
compile_args["weighted_metrics"]
|
|
)
|
|
clone.compile(**compile_args)
|
|
|
|
return clone
|