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)