def test_arg_and_kwarg_mix(self):
input_layer = DummyLayer()
input_layer_2 = DummyLayer()
a = DummyTensor()
node_a = node_module.Node(layer=input_layer, outputs=a)
b = DummyTensor()
node_b = node_module.Node(layer=input_layer_2, outputs=b)
arg_2 = DummyTensor()
arg_3 = DummyTensor()
node_c = node_module.Node(layer=input_layer, outputs=arg_3)
kwarg_x = DummyTensor()
kwarg_y = DummyTensor()
node_d = node_module.Node(layer=input_layer, outputs=kwarg_y)
merge_layer = DummyLayer()
merged = DummyTensor()
node = node_module.Node(
layer=merge_layer,
call_args=([a, b], arg_2, arg_3),
call_kwargs={
"x": kwarg_x,
"y": kwarg_y
},
outputs=merged,
)
(
merge_layer,
merge_node_index,
merge_tensor_index,
) = merged._keras_history
# Check the saved call args/kwargs
self.assertEqual(([a, b], arg_2, arg_3), node.call_args)
self.assertEqual({"x": kwarg_x, "y": kwarg_y}, node.call_kwargs)
# Only the inputs that were produced by input nodes should appear in
# keras_tensors
self.assertEqual({a, b, arg_3, kwarg_y}, set(node.keras_inputs))
self.assertEqual(set(node.parent_nodes),
{node_a, node_b, node_c, node_d})
# Check the layer wirings
self.assertEqual(merge_node_index, 0)
self.assertEqual(merge_tensor_index, 0)
self.assertLen(merge_layer._inbound_nodes, 1)
self.assertLen(merge_layer._outbound_nodes, 0)
self.assertLen(input_layer._outbound_nodes, 3)
self.assertLen(input_layer_2._outbound_nodes, 1)
self.assertLen(merge_layer._inbound_nodes[0].input_tensors, 2)
self.assertEqual(merge_layer._inbound_nodes[0].input_tensors, [a, b])
self.assertLen(merge_layer._inbound_nodes[0].inbound_layers, 4)
def clone_graph_nodes(inputs, outputs):
"""Clone the `Node` between the inputs and output tensors.
This function is used to create a new functional model from any intermediate
keras tensors. The clone of the nodes mimic the behavior of reconstructing
the functional graph network by re-executing all the __call__ methods. The
cloned nodes will be appended to the layers.
Note that a new tf.keras.Inputs will be created for any items in the
`inputs`
Args:
inputs: A nested structure of keras_tensors.
outputs: A nested structure of keras_tensors.
Returns:
A pair of inputs and outputs, with cloned keras_tensors. They can be used
to create a new functional model.
"""
nodes_to_clone = find_nodes_by_inputs_and_outputs(inputs, outputs)
cloned_inputs = []
cloned_outputs = []
# We not only need to create copies of Nodes (mimic the calls), also need to
# clone keras_tensors to avoid the override of _keras_history attached on
# the keras_tensor. The following dict is used to track any keras tensor we
# cloned The key is the string ID of the original keras tensor, and value is
# the cloned keras_tensor instance.
kt_id_mapping = {}
for kt_input in tf.nest.flatten(inputs):
if kt_input.node.is_input:
# For any existing keras_tensor from tf.keras.Input, we leave them
# as is.
cloned_inputs.append(kt_input)
kt_id_mapping[id(kt_input)] = kt_input
else:
# We need to create a new tf.keras.Input for any intermediate
# keras_tensor
cpy = _clone_keras_tensor(kt_input)
cloned_input = input_layer_module.Input(tensor=cpy)
cloned_inputs.append(cloned_input)
kt_id_mapping[id(kt_input)] = cloned_input
cloned_inputs = tf.nest.pack_sequence_as(inputs, cloned_inputs)
for kt_output in tf.nest.flatten(outputs):
cpy = _clone_keras_tensor(kt_output)
# We reuse the _keras_history here, which contains the old information.
# It is used in the Node constructor to check if the tensor
# "is_keras_tensor()" The history will be override by the Node
# constructor anyway for the corresponding layer output anyway.
cpy._keras_history = kt_output._keras_history
cloned_outputs.append(cpy)
kt_id_mapping[id(kt_output)] = cpy
cloned_outputs = tf.nest.pack_sequence_as(outputs, cloned_outputs)
for node in nodes_to_clone:
# Clone any keras_tensors to avoid override of _keras_history
# Or reuse an existing keras_tensor if it has already been cloned.
output_copy = clone_keras_tensors(node.output_tensors, kt_id_mapping)
call_args_copy = clone_keras_tensors(node.call_args, kt_id_mapping)
call_kwargs_copy = clone_keras_tensors(node.call_kwargs, kt_id_mapping)
# Creating new nodes based on the existing node information. Node wires
# itself to inbound and outbound layers. The Node constructor actually
# updates this layer's self._inbound_nodes, sets _keras_history on the
# outputs, and adds itself to the `_outbound_nodes` of the layers that
# produced the inputs to this layer call.
node_module.Node(
node.layer,
call_args=call_args_copy,
call_kwargs=call_kwargs_copy,
outputs=output_copy,
)
return cloned_inputs, cloned_outputs
def __init__(
self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=None,
name=None,
ragged=None,
type_spec=None,
**kwargs,
):
self._init_input_shape = input_shape
self._init_batch_size = batch_size
self._init_dtype = dtype
self._init_sparse = sparse
self._init_ragged = ragged
self._init_type_spec = type_spec
strategy = tf.distribute.get_strategy()
if (strategy and batch_size is not None
and distributed_training_utils.global_batch_size_supported(
strategy)):
if batch_size % strategy.num_replicas_in_sync != 0:
raise ValueError(
"The `batch_size` argument ({}) must be divisible by "
"the number of replicas ({})".format(
batch_size, strategy.num_replicas_in_sync))
batch_size = batch_size // strategy.num_replicas_in_sync
if "batch_input_shape" in kwargs:
batch_input_shape = kwargs.pop("batch_input_shape")
if input_shape and batch_input_shape:
raise ValueError("Only provide the input_shape OR "
"batch_input_shape argument to "
"InputLayer, not both at the same time.")
# Set the input shape and batch size from the batch_input_shape.
# Note that batch_input_shape can be None (unknown rank) or [] (scalar),
# in which case the batch size must be None.
if batch_input_shape:
batch_size = batch_input_shape[0]
input_shape = batch_input_shape[1:]
if kwargs:
raise ValueError(
f"Unrecognized keyword arguments: {list(kwargs.keys())}")
if sparse and ragged:
raise ValueError(
"Cannot set both sparse and ragged to True in a Keras input.")
if not name:
prefix = "input"
name = prefix + "_" + str(backend.get_uid(prefix))
if not dtype:
if input_tensor is None:
dtype = backend.floatx()
else:
dtype = backend.dtype(input_tensor)
elif input_tensor is not None and input_tensor.dtype != dtype:
raise ValueError(
"`input_tensor.dtype` differs from `dtype`. Received: "
f"input_tensor.dtype={input_tensor.dtype} "
f"but expected dtype={dtype}")
super().__init__(dtype=dtype, name=name)
self.built = True
self.sparse = True if sparse else False
self.ragged = True if ragged else False
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tf.TensorShape):
input_shape = tuple(input_shape.as_list())
elif isinstance(input_shape, int):
input_shape = (input_shape, )
if type_spec is not None:
args_that_must_be_none = [
("(input_)shape", self._init_input_shape),
("batch_size", self._init_batch_size),
("dtype", self._init_dtype),
("input_tensor", input_tensor),
("sparse", self._init_sparse),
("ragged", self._init_ragged),
]
for arg_name, arg in args_that_must_be_none:
_assert_other_arg_none(arg_name, arg)
if not tf.compat.v1.executing_eagerly_outside_functions():
raise ValueError(
"Creating Keras inputs from a type_spec is only "
"supported when eager execution is enabled.")
input_tensor = keras_tensor.keras_tensor_from_type_spec(type_spec)
if isinstance(input_tensor, keras_tensor.SparseKerasTensor):
self.sparse = True
if isinstance(input_tensor, keras_tensor.RaggedKerasTensor):
self.ragged = True
self.is_placeholder = True
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
elif input_tensor is None:
if input_shape is not None:
batch_input_shape = (batch_size, ) + tuple(input_shape)
else:
batch_input_shape = None
graph = backend.get_graph()
with graph.as_default():
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=sparse,
ragged=ragged,
)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
else:
if tf.compat.v1.executing_eagerly_outside_functions():
if not isinstance(input_tensor, keras_tensor.KerasTensor):
input_tensor = keras_tensor.keras_tensor_from_tensor(
input_tensor)
else:
if not tf_utils.is_symbolic_tensor(input_tensor):
raise ValueError(
"You should not pass an EagerTensor to `Input`. "
"For example, instead of creating an "
"`InputLayer`, you should instantiate your model "
"and directly call it on your input.")
self.is_placeholder = False
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
# Create an input node.
input_tensor._keras_mask = None
node_module.Node(layer=self, outputs=input_tensor)
# Store type spec
if isinstance(input_tensor, keras_tensor.KerasTensor) or (
tf_utils.is_extension_type(input_tensor)):
self._type_spec = (input_tensor._type_spec) # pylint: disable=protected-access
else:
self._type_spec = tf.TensorSpec(
shape=input_tensor.shape,
dtype=input_tensor.dtype,
name=self.name,
)
def __init__(self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=None,
name=None,
ragged=None,
type_spec=None,
**kwargs):
self._init_input_shape = input_shape
self._init_batch_size = batch_size
self._init_dtype = dtype
self._init_sparse = sparse
self._init_ragged = ragged
self._init_type_spec = type_spec
strategy = tf.distribute.get_strategy()
if strategy and batch_size is not None and \
distributed_training_utils.global_batch_size_supported(strategy):
if batch_size % strategy.num_replicas_in_sync != 0:
raise ValueError(
'The `batch_size` argument ({}) must be divisible by '
'the number of replicas ({})'.format(
batch_size, strategy.num_replicas_in_sync))
batch_size = batch_size // strategy.num_replicas_in_sync
if 'batch_input_shape' in kwargs:
batch_input_shape = kwargs.pop('batch_input_shape')
if input_shape and batch_input_shape:
raise ValueError('Only provide the input_shape OR '
'batch_input_shape argument to '
'InputLayer, not both at the same time.')
batch_size = batch_input_shape[0]
input_shape = batch_input_shape[1:]
if kwargs:
raise ValueError('Unrecognized keyword arguments:', kwargs.keys())
if sparse and ragged:
raise ValueError(
'Cannot set both sparse and ragged to True in a Keras input.')
if not name:
prefix = 'input'
name = prefix + '_' + str(backend.get_uid(prefix))
if not dtype:
if input_tensor is None:
dtype = backend.floatx()
else:
dtype = backend.dtype(input_tensor)
elif input_tensor is not None and input_tensor.dtype != dtype:
raise ValueError(
'`input_tensor.dtype` differs from `dtype`: %s vs. %s' %
(input_tensor.dtype, dtype))
super(InputLayer, self).__init__(dtype=dtype, name=name)
self.built = True
self.sparse = True if sparse else False
self.ragged = True if ragged else False
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tf.TensorShape):
input_shape = tuple(input_shape.as_list())
elif isinstance(input_shape, int):
input_shape = (input_shape, )
if type_spec is not None:
args_that_must_be_none = [
('(input_)shape', self._init_input_shape),
('batch_size', self._init_batch_size),
('dtype', self._init_dtype),
('input_tensor', input_tensor),
('sparse', self._init_sparse),
('ragged', self._init_ragged),
]
for arg_name, arg in args_that_must_be_none:
_assert_other_arg_none(arg_name, arg)
if not keras_tensor.keras_tensors_enabled():
raise ValueError(
'Creating Keras inputs from a type_spec is only '
'supported when eager execution is enabled.')
input_tensor = keras_tensor.keras_tensor_from_type_spec(type_spec)
if isinstance(input_tensor, keras_tensor.SparseKerasTensor):
self.sparse = True
if isinstance(input_tensor, keras_tensor.RaggedKerasTensor):
self.ragged = True
self.is_placeholder = True
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
elif input_tensor is None:
if input_shape is not None:
batch_input_shape = (batch_size, ) + tuple(input_shape)
else:
batch_input_shape = None
graph = backend.get_graph()
with graph.as_default():
input_tensor = backend.placeholder(shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=sparse,
ragged=ragged)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
else:
if keras_tensor.keras_tensors_enabled():
if not isinstance(input_tensor, keras_tensor.KerasTensor):
input_tensor = keras_tensor.keras_tensor_from_tensor(
input_tensor)
else:
if not tf_utils.is_symbolic_tensor(input_tensor):
raise ValueError(
'You should not pass an EagerTensor to `Input`. '
'For example, instead of creating an '
'InputLayer, you should instantiate your model and '
'directly call it on your input.')
self.is_placeholder = False
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
# Create an input node.
input_tensor._keras_mask = None
node_module.Node(layer=self, outputs=input_tensor)
# Store type spec
if isinstance(input_tensor, keras_tensor.KerasTensor) or (
tf_utils.is_extension_type(input_tensor)):
self._type_spec = input_tensor._type_spec # pylint: disable=protected-access
else:
self._type_spec = tf.TensorSpec(shape=input_tensor.shape,
dtype=input_tensor.dtype,
name=self.name)
