def call(self,
inputs,
mask=None,
training=None,
initial_state=None,
constants=None):
# The input should be dense, padded with zeros. If a ragged input is fed
# into the layer, it is padded and the row lengths are used for masking.
inputs, row_lengths = backend.convert_inputs_if_ragged(inputs)
is_ragged_input = (row_lengths is not None)
self._validate_args_if_ragged(is_ragged_input, mask)
inputs, initial_state, constants = self._process_inputs(
inputs, initial_state, constants)
self._maybe_reset_cell_dropout_mask(self.cell)
if isinstance(self.cell, StackedRNNCells):
for cell in self.cell.cells:
self._maybe_reset_cell_dropout_mask(cell)
if mask is not None:
# Time step masks must be the same for each input.
# TODO(scottzhu): Should we accept multiple different masks?
mask = tf.nest.flatten(mask)[0]
if tf.nest.is_nested(inputs):
# In the case of nested input, use the first element for shape check.
input_shape = backend.int_shape(tf.nest.flatten(inputs)[0])
else:
input_shape = backend.int_shape(inputs)
timesteps = input_shape[0] if self.time_major else input_shape[1]
if self.unroll and timesteps is None:
raise ValueError('Cannot unroll a RNN if the '
'time dimension is undefined. \n'
'- If using a Sequential model, '
'specify the time dimension by passing '
'an `input_shape` or `batch_input_shape` '
'argument to your first layer. If your '
'first layer is an Embedding, you can '
'also use the `input_length` argument.\n'
'- If using the functional API, specify '
'the time dimension by passing a `shape` '
'or `batch_shape` argument to your Input layer.')
kwargs = {}
if generic_utils.has_arg(self.cell.call, 'training'):
kwargs['training'] = training
# TF RNN cells expect single tensor as state instead of list wrapped tensor.
is_tf_rnn_cell = getattr(self.cell, '_is_tf_rnn_cell',
None) is not None
# Use the __call__ function for callable objects, eg layers, so that it
# will have the proper name scopes for the ops, etc.
cell_call_fn = self.cell.__call__ if callable(
self.cell) else self.cell.call
if constants:
if not generic_utils.has_arg(self.cell.call, 'constants'):
raise ValueError(
f'RNN cell {self.cell} does not support constants. '
f'Received: constants={constants}')
def step(inputs, states):
constants = states[-self._num_constants:] # pylint: disable=invalid-unary-operand-type
states = states[:-self._num_constants] # pylint: disable=invalid-unary-operand-type
states = states[0] if len(
states) == 1 and is_tf_rnn_cell else states
output, new_states = cell_call_fn(inputs,
states,
constants=constants,
**kwargs)
if not tf.nest.is_nested(new_states):
new_states = [new_states]
return output, new_states
else:
def step(inputs, states):
states = states[0] if len(
states) == 1 and is_tf_rnn_cell else states
output, new_states = cell_call_fn(inputs, states, **kwargs)
if not tf.nest.is_nested(new_states):
new_states = [new_states]
return output, new_states
last_output, outputs, states = backend.rnn(
step,
inputs,
initial_state,
constants=constants,
go_backwards=self.go_backwards,
mask=mask,
unroll=self.unroll,
input_length=row_lengths if row_lengths is not None else timesteps,
time_major=self.time_major,
zero_output_for_mask=self.zero_output_for_mask)
if self.stateful:
updates = [
tf.compat.v1.assign(self_state,
tf.cast(state, self_state.dtype))
for self_state, state in zip(tf.nest.flatten(self.states),
tf.nest.flatten(states))
]
self.add_update(updates)
if self.return_sequences:
output = backend.maybe_convert_to_ragged(
is_ragged_input,
outputs,
row_lengths,
go_backwards=self.go_backwards)
else:
output = last_output
if self.return_state:
if not isinstance(states, (list, tuple)):
states = [states]
else:
states = list(states)
return generic_utils.to_list(output) + states
else:
return output
def call(self, inputs, mask=None, training=None, initial_state=None):
# The input should be dense, padded with zeros. If a ragged input is fed
# into the layer, it is padded and the row lengths are used for masking.
inputs, row_lengths = backend.convert_inputs_if_ragged(inputs)
is_ragged_input = row_lengths is not None
self._validate_args_if_ragged(is_ragged_input, mask)
# GRU does not support constants. Ignore it during process.
inputs, initial_state, _ = self._process_inputs(
inputs, initial_state, None)
if isinstance(mask, list):
mask = mask[0]
input_shape = backend.int_shape(inputs)
timesteps = input_shape[0] if self.time_major else input_shape[1]
if not self._could_use_gpu_kernel:
kwargs = {"training": training}
self._maybe_reset_cell_dropout_mask(self.cell)
def step(cell_inputs, cell_states):
return self.cell(cell_inputs, cell_states, **kwargs)
last_output, outputs, states = backend.rnn(
step,
inputs,
initial_state,
constants=None,
go_backwards=self.go_backwards,
mask=mask,
unroll=self.unroll,
input_length=row_lengths
if row_lengths is not None else timesteps,
time_major=self.time_major,
zero_output_for_mask=self.zero_output_for_mask,
return_all_outputs=self.return_sequences,
)
# This is a dummy tensor for testing purpose.
runtime = gru_lstm_utils.runtime(gru_lstm_utils.RUNTIME_UNKNOWN)
else:
last_output, outputs, runtime, states = self._defun_gru_call(
inputs, initial_state, training, mask, row_lengths)
if self.stateful:
updates = [
tf.compat.v1.assign(self.states[0],
tf.cast(states[0], self.states[0].dtype))
]
self.add_update(updates)
if self.return_sequences:
output = backend.maybe_convert_to_ragged(
is_ragged_input,
outputs,
row_lengths,
go_backwards=self.go_backwards,
)
else:
output = last_output
if self.return_state:
return [output] + list(states)
elif self._return_runtime:
return output, runtime
else:
return output
def call(self, inputs, training=None, mask=None):
kwargs = {}
if generic_utils.has_arg(self.layer.call, 'training'):
kwargs['training'] = training
input_shape = tf.nest.map_structure(
lambda x: tf.TensorShape(K.int_shape(x)), inputs)
batch_size = tf_utils.convert_shapes(input_shape)
batch_size = tf.nest.flatten(batch_size)[0]
if batch_size and not self._always_use_reshape:
inputs, row_lengths = K.convert_inputs_if_ragged(inputs)
is_ragged_input = row_lengths is not None
input_length = tf_utils.convert_shapes(input_shape)
input_length = tf.nest.flatten(input_length)[1]
# batch size matters, use rnn-based implementation
def step(x, _):
output = self.layer(x, **kwargs)
return output, []
_, outputs, _ = K.rnn(
step,
inputs,
initial_states=[],
input_length=row_lengths[0] if is_ragged_input else input_length,
mask=mask,
unroll=False)
# pylint: disable=g-long-lambda
y = tf.nest.map_structure(
lambda output: K.maybe_convert_to_ragged(is_ragged_input, output,
row_lengths), outputs)
else:
# No batch size specified, therefore the layer will be able
# to process batches of any size.
# We can go with reshape-based implementation for performance.
is_ragged_input = tf.nest.map_structure(
lambda x: isinstance(x, tf.RaggedTensor), inputs)
is_ragged_input = tf.nest.flatten(is_ragged_input)
if all(is_ragged_input):
input_values = tf.nest.map_structure(lambda x: x.values, inputs)
input_row_lenghts = tf.nest.map_structure(
lambda x: x.nested_row_lengths()[0], inputs)
y = self.layer(input_values, **kwargs)
y = tf.nest.map_structure(tf.RaggedTensor.from_row_lengths, y,
input_row_lenghts)
elif any(is_ragged_input):
raise ValueError('All inputs has to be either ragged or not, '
'but not mixed. You passed: {}'.format(inputs))
else:
input_length = tf_utils.convert_shapes(input_shape)
input_length = tf.nest.flatten(input_length)[1]
if not input_length:
input_length = tf.nest.map_structure(lambda x: tf.compat.v1.shape(x)[1],
inputs)
input_length = generic_utils.to_list(tf.nest.flatten(input_length))[0]
inner_input_shape = tf.nest.map_structure(
lambda x: self._get_shape_tuple((-1,), x, 2), inputs)
# Shape: (num_samples * timesteps, ...). And track the
# transformation in self._input_map.
inputs = tf.__internal__.nest.map_structure_up_to(inputs, tf.reshape, inputs,
inner_input_shape)
# (num_samples * timesteps, ...)
if generic_utils.has_arg(self.layer.call, 'mask') and mask is not None:
inner_mask_shape = self._get_shape_tuple((-1,), mask, 2)
kwargs['mask'] = K.reshape(mask, inner_mask_shape)
y = self.layer(inputs, **kwargs)
# Shape: (num_samples, timesteps, ...)
output_shape = self.compute_output_shape(input_shape)
# pylint: disable=g-long-lambda
output_shape = tf.nest.map_structure(
lambda tensor, int_shape: self._get_shape_tuple(
(-1, input_length), tensor, 1, int_shape[2:]), y, output_shape)
y = tf.__internal__.nest.map_structure_up_to(y, tf.reshape, y, output_shape)
if not tf.executing_eagerly():
# Set the static shape for the result since it might be lost during
# array_ops reshape, eg, some `None` dim in the result could be
# inferred.
tf.__internal__.nest.map_structure_up_to(
y, lambda tensor, shape: tensor.set_shape(shape), y,
self.compute_output_shape(input_shape))
return y
def call(self, inputs, mask=None, training=None, initial_state=None):
# The input should be dense, padded with zeros. If a ragged input is fed
# into the layer, it is padded and the row lengths are used for masking.
inputs, row_lengths = backend.convert_inputs_if_ragged(inputs)
is_ragged_input = (row_lengths is not None)
self._validate_args_if_ragged(is_ragged_input, mask)
# LSTM does not support constants. Ignore it during process.
inputs, initial_state, _ = self._process_inputs(inputs, initial_state, None)
if isinstance(mask, list):
mask = mask[0]
input_shape = backend.int_shape(inputs)
timesteps = input_shape[0] if self.time_major else input_shape[1]
if not self._could_use_gpu_kernel:
# Fall back to use the normal LSTM.
kwargs = {'training': training}
self._maybe_reset_cell_dropout_mask(self.cell)
def step(inputs, states):
return self.cell(inputs, states, **kwargs)
last_output, outputs, states = backend.rnn(
step,
inputs,
initial_state,
constants=None,
go_backwards=self.go_backwards,
mask=mask,
unroll=self.unroll,
input_length=row_lengths if row_lengths is not None else timesteps,
time_major=self.time_major,
zero_output_for_mask=self.zero_output_for_mask)
runtime = gru_lstm_utils.runtime(gru_lstm_utils.RUNTIME_UNKNOWN)
else:
# Use the new defun approach for backend implementation swap.
# Note that different implementations need to have same function
# signature, eg, the tensor parameters need to have same shape and dtypes.
# Since the cuDNN has an extra set of bias, those bias will be passed to
# both normal and cuDNN implementations.
self.reset_dropout_mask()
dropout_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
if dropout_mask is not None:
inputs = inputs * dropout_mask[0]
if gru_lstm_utils.use_new_gru_lstm_impl():
lstm_kwargs = {
'inputs':
inputs,
'init_h':
gru_lstm_utils.read_variable_value(initial_state[0]),
'init_c':
gru_lstm_utils.read_variable_value(initial_state[1]),
'kernel':
gru_lstm_utils.read_variable_value(self.cell.kernel),
'recurrent_kernel':
gru_lstm_utils.read_variable_value(self.cell.recurrent_kernel),
'bias':
gru_lstm_utils.read_variable_value(self.cell.bias),
'mask':
mask,
'time_major':
self.time_major,
'go_backwards':
self.go_backwards,
'sequence_lengths':
row_lengths,
'zero_output_for_mask':
self.zero_output_for_mask,
}
(last_output, outputs, new_h, new_c,
runtime) = self._defun_wrapper.defun_layer(**lstm_kwargs)
else:
gpu_lstm_kwargs = {
'inputs':
inputs,
'init_h':
gru_lstm_utils.read_variable_value(initial_state[0]),
'init_c':
gru_lstm_utils.read_variable_value(initial_state[1]),
'kernel':
gru_lstm_utils.read_variable_value(self.cell.kernel),
'recurrent_kernel':
gru_lstm_utils.read_variable_value(self.cell.recurrent_kernel),
'bias':
gru_lstm_utils.read_variable_value(self.cell.bias),
'mask':
mask,
'time_major':
self.time_major,
'go_backwards':
self.go_backwards,
'sequence_lengths':
row_lengths
}
normal_lstm_kwargs = gpu_lstm_kwargs.copy()
normal_lstm_kwargs.update({
'zero_output_for_mask': self.zero_output_for_mask,
})
if tf.executing_eagerly():
device_type = gru_lstm_utils.get_context_device_type()
can_use_gpu = (
# Either user specified GPU or unspecified but GPU is available.
(device_type == gru_lstm_utils.GPU_DEVICE_NAME or
(device_type is None
and tf.config.list_logical_devices('GPU'))) and
(mask is None or
gru_lstm_utils.is_cudnn_supported_inputs(mask, self.time_major)))
# Under eager context, check the device placement and prefer the
# GPU implementation when GPU is available.
if can_use_gpu:
last_output, outputs, new_h, new_c, runtime = gpu_lstm(
**gpu_lstm_kwargs)
else:
last_output, outputs, new_h, new_c, runtime = standard_lstm(
**normal_lstm_kwargs)
else:
(last_output, outputs, new_h, new_c,
runtime) = lstm_with_backend_selection(**normal_lstm_kwargs)
states = [new_h, new_c]
if self.stateful:
updates = [
tf.compat.v1.assign(self_state, tf.cast(state, self_state.dtype))
for self_state, state in zip(self.states, states)
]
self.add_update(updates)
if self.return_sequences:
output = backend.maybe_convert_to_ragged(
is_ragged_input, outputs, row_lengths, go_backwards=self.go_backwards)
else:
output = last_output
if self.return_state:
return [output] + list(states)
elif self.return_runtime:
return output, runtime
else:
return output
