345 lines
14 KiB
Python
345 lines
14 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.
|
||
|
|
# ==============================================================================
|
||
|
|
|
||
|
|
|
||
|
|
"""Contains the `Node` class."""
|
||
|
|
|
||
|
|
import collections
|
||
|
|
import copy
|
||
|
|
import json
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import tensorflow.compat.v2 as tf
|
||
|
|
|
||
|
|
from keras import backend
|
||
|
|
from keras.engine import base_layer_utils
|
||
|
|
from keras.saving.legacy.saved_model import json_utils
|
||
|
|
from keras.utils import tf_utils
|
||
|
|
|
||
|
|
_CONSTANT_VALUE = "_CONSTANT_VALUE"
|
||
|
|
# Using dict to avoid conflict with constant string tensor.
|
||
|
|
_COMPOSITE_TYPE = {"_TYPE": "COMPOSITE"}
|
||
|
|
|
||
|
|
|
||
|
|
class Node:
|
||
|
|
"""A `Node` describes a layer `__call__()` event.
|
||
|
|
|
||
|
|
A Functional model is a DAG with `Node` instances as nodes, and
|
||
|
|
`KerasTensor` instances as edges. Nodes aren't `Layer` instances, because a
|
||
|
|
single layer could be called multiple times, which would result in graph
|
||
|
|
cycles.
|
||
|
|
|
||
|
|
A `__call__()` event involves input tensors (and other input arguments),
|
||
|
|
the layer that was called, and the resulting output tensors.
|
||
|
|
A `Node` will include all this information.
|
||
|
|
|
||
|
|
Since a single `Layer` could be called multiple times, the `Node` instances
|
||
|
|
are stored on layers as a list. Each time a layer is called a node is added
|
||
|
|
to `layer._inbound_nodes`. Each time the output of a layer is used by
|
||
|
|
another layer, a node is added to `layer._outbound_nodes`.
|
||
|
|
|
||
|
|
Every `KerasTensor` instance has a `KerasHistory` object attached,
|
||
|
|
which tracks the `Node` that records the `__call__()` event that created
|
||
|
|
the tensor. By recursively walking through `Node` instances
|
||
|
|
via the `KerasHistory` metadata of `KerasTensor` instances, once can
|
||
|
|
retrieve the entire DAG of a Functional model.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
layer: The layer that was called in the `Layer.__call__()`
|
||
|
|
event that this node represents.
|
||
|
|
call_args: The positional arguments the layer was called with.
|
||
|
|
call_kwargs: The keyword arguments the layer was called with.
|
||
|
|
outputs: The output tensors of the `Layer.__call__()`
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, layer, call_args=None, call_kwargs=None, outputs=None):
|
||
|
|
call_args = [] if call_args is None else call_args
|
||
|
|
call_kwargs = {} if call_kwargs is None else call_kwargs
|
||
|
|
outputs = [] if outputs is None else outputs
|
||
|
|
|
||
|
|
self.layer = layer
|
||
|
|
self.is_input = not call_args and not call_kwargs
|
||
|
|
|
||
|
|
# These arguments are user-provided. Copy the structures here so that
|
||
|
|
# future user modifications do not affect the node's metadata.
|
||
|
|
# We copy using map_structure rather than python's shallow or deep copy,
|
||
|
|
# because the args can be data structures (so shallow copy is
|
||
|
|
# insufficient), but individual values might not support copy.copy
|
||
|
|
# or be too expensive to deep copy.
|
||
|
|
call_args = tf.nest.map_structure(lambda t: t, call_args)
|
||
|
|
call_kwargs = tf.nest.map_structure(lambda t: t, call_kwargs)
|
||
|
|
self.outputs = tf.nest.map_structure(lambda t: t, outputs)
|
||
|
|
self.call_args = call_args
|
||
|
|
self.call_kwargs = call_kwargs
|
||
|
|
|
||
|
|
# Cached for performance.
|
||
|
|
self._flat_arguments = tf.nest.flatten(
|
||
|
|
(self.call_args, self.call_kwargs)
|
||
|
|
)
|
||
|
|
# Used to avoid expensive `nest` operations in the most common case.
|
||
|
|
self._single_positional_tensor_passed = (
|
||
|
|
not self.call_kwargs
|
||
|
|
and len(self.call_args) == 1
|
||
|
|
and tf.is_tensor(self.call_args[0])
|
||
|
|
)
|
||
|
|
|
||
|
|
if not tf.compat.v1.executing_eagerly_outside_functions():
|
||
|
|
# Create TensorFlowOpLayers if needed (in TF1)
|
||
|
|
for obj in self._flat_arguments:
|
||
|
|
if isinstance(
|
||
|
|
obj, tf.Tensor
|
||
|
|
) and base_layer_utils.needs_keras_history(
|
||
|
|
obj, ignore_call_context=True
|
||
|
|
):
|
||
|
|
base_layer_utils.create_keras_history(obj)
|
||
|
|
|
||
|
|
self._keras_inputs = []
|
||
|
|
self._keras_inputs_ids_and_indices = []
|
||
|
|
for i, ele in enumerate(self._flat_arguments):
|
||
|
|
if is_keras_tensor(ele):
|
||
|
|
self._keras_inputs.append(ele)
|
||
|
|
kt_id = str(id(ele))
|
||
|
|
kt_index = i
|
||
|
|
self._keras_inputs_ids_and_indices.append((kt_id, kt_index))
|
||
|
|
|
||
|
|
# Wire up Node to Layers.
|
||
|
|
self.layer._inbound_nodes.append(self)
|
||
|
|
for kt in self.keras_inputs:
|
||
|
|
inbound_layer = kt._keras_history.layer
|
||
|
|
if inbound_layer is not None: # `None` for `Input` tensors.
|
||
|
|
inbound_layer._outbound_nodes.append(self)
|
||
|
|
|
||
|
|
# Set metadata on outputs.
|
||
|
|
node_index = len(self.layer._inbound_nodes) - 1
|
||
|
|
for i, tensor in enumerate(tf.nest.flatten(outputs)):
|
||
|
|
tensor._keras_history = KerasHistory(
|
||
|
|
layer=layer, node_index=node_index, tensor_index=i
|
||
|
|
)
|
||
|
|
|
||
|
|
# Cached for performance.
|
||
|
|
self.flat_input_ids = [str(id(t)) for t in self._keras_inputs]
|
||
|
|
self.flat_output_ids = [
|
||
|
|
str(id(t)) for t in tf.nest.flatten(self.outputs)
|
||
|
|
]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def keras_inputs(self):
|
||
|
|
"""Tensors input to this node that can be traced back to a
|
||
|
|
`keras.Input`."""
|
||
|
|
return self._keras_inputs
|
||
|
|
|
||
|
|
@property
|
||
|
|
def parent_nodes(self):
|
||
|
|
"""Returns all the `Node`s whose output this node immediately depends
|
||
|
|
on."""
|
||
|
|
node_deps = []
|
||
|
|
for kt in self.keras_inputs:
|
||
|
|
layer = kt._keras_history.layer
|
||
|
|
node_index = kt._keras_history.node_index
|
||
|
|
if layer is not None: # `None` for `Input` tensors.
|
||
|
|
node_deps.append(layer._inbound_nodes[node_index])
|
||
|
|
return node_deps
|
||
|
|
|
||
|
|
def iterate_inbound(self):
|
||
|
|
"""Yields tuples representing the data inbound from other nodes.
|
||
|
|
|
||
|
|
Yields:
|
||
|
|
tuples like: (inbound_layer, node_index, tensor_index, tensor).
|
||
|
|
"""
|
||
|
|
for kt in self.keras_inputs:
|
||
|
|
keras_history = kt._keras_history
|
||
|
|
layer = keras_history.layer
|
||
|
|
node_index = keras_history.node_index
|
||
|
|
tensor_index = keras_history.tensor_index
|
||
|
|
yield layer, node_index, tensor_index, kt
|
||
|
|
|
||
|
|
def map_arguments(self, tensor_dict):
|
||
|
|
"""Maps Keras Tensors to computed Tensors using `tensor_dict`."""
|
||
|
|
if self._single_positional_tensor_passed:
|
||
|
|
# Performance optimization for most common case.
|
||
|
|
kt_id, _ = self._keras_inputs_ids_and_indices[0]
|
||
|
|
return (tensor_dict[kt_id].pop(),), {}
|
||
|
|
else:
|
||
|
|
flat_arguments = copy.copy(self._flat_arguments)
|
||
|
|
for kt_id, kt_index in self._keras_inputs_ids_and_indices:
|
||
|
|
flat_arguments[kt_index] = tensor_dict[kt_id].pop()
|
||
|
|
|
||
|
|
args, kwargs = tf.nest.pack_sequence_as(
|
||
|
|
(self.call_args, self.call_kwargs), flat_arguments
|
||
|
|
)
|
||
|
|
return args, kwargs
|
||
|
|
|
||
|
|
def serialize(self, make_node_key, node_conversion_map):
|
||
|
|
"""Serializes `Node` for Functional API's `get_config`."""
|
||
|
|
# Serialization still special-cases first argument.
|
||
|
|
args, kwargs = self.call_args, self.call_kwargs
|
||
|
|
inputs, args, kwargs = self.layer._call_spec.split_out_first_arg(
|
||
|
|
args, kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
# Treat everything other than first argument as a kwarg.
|
||
|
|
arguments = dict(zip(self.layer._call_spec.arg_names[1:], args))
|
||
|
|
arguments.update(kwargs)
|
||
|
|
kwargs = arguments
|
||
|
|
|
||
|
|
def _serialize_keras_tensor(t):
|
||
|
|
"""Serializes a single Tensor passed to `call`."""
|
||
|
|
if hasattr(t, "_keras_history"):
|
||
|
|
kh = t._keras_history
|
||
|
|
node_index = kh.node_index
|
||
|
|
node_key = make_node_key(kh.layer.name, node_index)
|
||
|
|
new_node_index = node_conversion_map.get(node_key, 0)
|
||
|
|
return [kh.layer.name, new_node_index, kh.tensor_index]
|
||
|
|
|
||
|
|
if isinstance(t, np.ndarray):
|
||
|
|
return t.tolist()
|
||
|
|
|
||
|
|
if isinstance(t, tf.Tensor):
|
||
|
|
return backend.get_value(t).tolist()
|
||
|
|
|
||
|
|
# Not using json_utils to serialize both constant Tensor and
|
||
|
|
# constant CompositeTensor for saving format backward compatibility.
|
||
|
|
if isinstance(t, tf.__internal__.CompositeTensor):
|
||
|
|
return (_COMPOSITE_TYPE, json_utils.Encoder().encode(t))
|
||
|
|
|
||
|
|
return t
|
||
|
|
|
||
|
|
kwargs = tf.nest.map_structure(_serialize_keras_tensor, kwargs)
|
||
|
|
try:
|
||
|
|
json.dumps(kwargs, default=json_utils.get_json_type)
|
||
|
|
except TypeError:
|
||
|
|
kwarg_types = tf.nest.map_structure(type, kwargs)
|
||
|
|
raise TypeError(
|
||
|
|
"Layer "
|
||
|
|
+ self.layer.name
|
||
|
|
+ " was passed non-JSON-serializable arguments. "
|
||
|
|
+ "Arguments had types: "
|
||
|
|
+ str(kwarg_types)
|
||
|
|
+ ". They cannot be serialized out when saving the model."
|
||
|
|
)
|
||
|
|
|
||
|
|
# `kwargs` is added to each Tensor in the first arg. This should be
|
||
|
|
# changed in a future version of the serialization format.
|
||
|
|
def serialize_first_arg_tensor(t):
|
||
|
|
if is_keras_tensor(t):
|
||
|
|
kh = t._keras_history
|
||
|
|
node_index = kh.node_index
|
||
|
|
node_key = make_node_key(kh.layer.name, node_index)
|
||
|
|
new_node_index = node_conversion_map.get(node_key, 0)
|
||
|
|
data = [kh.layer.name, new_node_index, kh.tensor_index, kwargs]
|
||
|
|
else:
|
||
|
|
# If an element in the first call argument did not originate as
|
||
|
|
# a keras tensor and is a constant value, we save it using the
|
||
|
|
# format ['_CONSTANT_VALUE', -1,
|
||
|
|
# serialized_tensor_or_python_constant] (potentially including
|
||
|
|
# serialized kwargs in an optional 4th argument).
|
||
|
|
data = [_CONSTANT_VALUE, -1, _serialize_keras_tensor(t), kwargs]
|
||
|
|
return tf_utils.ListWrapper(data)
|
||
|
|
|
||
|
|
data = tf.nest.map_structure(serialize_first_arg_tensor, inputs)
|
||
|
|
if (
|
||
|
|
not tf.nest.is_nested(data)
|
||
|
|
and not self.layer._preserve_input_structure_in_config
|
||
|
|
):
|
||
|
|
data = [data]
|
||
|
|
data = tf_utils.convert_inner_node_data(data)
|
||
|
|
return data
|
||
|
|
|
||
|
|
#############################################################
|
||
|
|
# Properties for Backwards compatibility.
|
||
|
|
# These only check the first input argument
|
||
|
|
# As nodes are internal, they may be removed in the future.
|
||
|
|
#############################################################
|
||
|
|
|
||
|
|
@property
|
||
|
|
def input_tensors(self):
|
||
|
|
if self.is_input:
|
||
|
|
return [self.outputs] # Used in `Layer.input`.
|
||
|
|
return self.call_args[0]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def output_tensors(self):
|
||
|
|
if self.is_input:
|
||
|
|
return [self.outputs] # Used in `Layer.input`.
|
||
|
|
return self.outputs
|
||
|
|
|
||
|
|
@property
|
||
|
|
def input_shapes(self):
|
||
|
|
input_shapes = tf.nest.map_structure(
|
||
|
|
backend.int_shape, self.input_tensors
|
||
|
|
)
|
||
|
|
if len(input_shapes) == 1 and not self.is_input:
|
||
|
|
return input_shapes[0]
|
||
|
|
return input_shapes
|
||
|
|
|
||
|
|
@property
|
||
|
|
def output_shapes(self):
|
||
|
|
return tf.nest.map_structure(backend.int_shape, self.output_tensors)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def outbound_layer(self):
|
||
|
|
return self.layer
|
||
|
|
|
||
|
|
@property
|
||
|
|
def inbound_layers(self):
|
||
|
|
"""Return all layers that feed into the current node."""
|
||
|
|
if self.is_input:
|
||
|
|
return []
|
||
|
|
tensor_call_args = [
|
||
|
|
x
|
||
|
|
for x in self._flat_arguments
|
||
|
|
if tf.is_tensor(x) and hasattr(x, "_keras_history")
|
||
|
|
]
|
||
|
|
inbound_layers = tf.nest.map_structure(
|
||
|
|
lambda t: t._keras_history.layer, tensor_call_args
|
||
|
|
)
|
||
|
|
if len(inbound_layers) == 1:
|
||
|
|
return inbound_layers[0]
|
||
|
|
return inbound_layers
|
||
|
|
|
||
|
|
|
||
|
|
class KerasHistory(
|
||
|
|
collections.namedtuple(
|
||
|
|
"KerasHistory", ["layer", "node_index", "tensor_index"]
|
||
|
|
)
|
||
|
|
):
|
||
|
|
"""Tracks the Layer call that created a Tensor, for Keras Graph Networks.
|
||
|
|
|
||
|
|
During construction of Keras Graph Networks, this metadata is added to
|
||
|
|
each Tensor produced as the output of a Layer, starting with an
|
||
|
|
`InputLayer`. This allows Keras to track how each Tensor was produced, and
|
||
|
|
this information is later retraced by the `keras.engine.Network` class to
|
||
|
|
reconstruct the Keras Graph Network.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
layer: The Layer that produced the Tensor.
|
||
|
|
node_index: The specific call to the Layer that produced this Tensor.
|
||
|
|
Layers can be called multiple times in order to share weights. A new
|
||
|
|
node is created every time a Layer is called. The corresponding node
|
||
|
|
that represents the call event that produced the Tensor can be found at
|
||
|
|
`layer._inbound_nodes[node_index]`.
|
||
|
|
tensor_index: The output index for this Tensor. Always zero if the Layer
|
||
|
|
that produced this Tensor only has one output. Nested structures of
|
||
|
|
Tensors are deterministically assigned an index via `nest.flatten`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Added to maintain memory and performance characteristics of `namedtuple`
|
||
|
|
# while subclassing.
|
||
|
|
__slots__ = ()
|
||
|
|
|
||
|
|
|
||
|
|
def is_keras_tensor(obj):
|
||
|
|
return hasattr(obj, "_keras_history")
|