Ragged Tensors of non-text data as input for LSTM - python-3.x

Question

I am learning about ragged tensors with applications for particle tracking. I have the following minimal example which reproduces a the error i keep experiencing.
import tensorflow as tf
from tensorflow.keras.layers import Dense, LSTM, Input, TimeDistributed
from tensorflow.keras.models import Sequential
n=10
data_n = 32
batch_size=8
window_length=8
splits = [n]*data_n
#### Create a ragged tensor with shape (32, None, 8)
t0 = tf.zeros([data_n * n, window_length])
t1 = tf.RaggedTensor.from_row_lengths(t0, splits)
max_seq = t1.bounding_shape()[-1]
#### Define Model
def create_model(batch_size, window_length, max_seq):
lstm_model = Sequential([
Input(
batch_shape=[batch_size, None, window_length],
batch_size=batch_size,
dtype=tf.float32,
ragged=True
),
LSTM(
max_seq,
return_sequences=True,
input_shape=(window_length, None)
),
TimeDistributed(Dense(units=1))
])
return lstm_model
lstm_model = create_model(batch_size=batch_size, window_length=window_length, max_seq=max_seq)
lstm_model(t1[0:8])
When i execute the above code i get the following error:
---------------------------------------------------------------------------
_FallbackException Traceback (most recent call last)
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/ops/gen_cudnn_rnn_ops.py in cudnn_rnnv3(input, input_h, input_c, params, sequence_lengths, rnn_mode, input_mode, direction, dropout, seed, seed2, num_proj, is_training, time_major, name)
1889 try:
-> 1890 _result = pywrap_tfe.TFE_Py_FastPathExecute(
1891 _ctx._context_handle, tld.device_name, "CudnnRNNV3", name,
_FallbackException: Expecting float value for attr dropout, got int
During handling of the above exception, another exception occurred:
InvalidArgumentError Traceback (most recent call last)
<ipython-input-19-7609e2877e20> in <module>
1 lstm_model = create_model(batch_size=batch_size, window_length=window_length, max_seq=max_seq)
----> 2 lstm_model(t1[0:8])
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
966 with base_layer_utils.autocast_context_manager(
967 self._compute_dtype):
--> 968 outputs = self.call(cast_inputs, *args, **kwargs)
969 self._handle_activity_regularization(inputs, outputs)
970 self._set_mask_metadata(inputs, outputs, input_masks)
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/sequential.py in call(self, inputs, training, mask)
275 if not self.built:
276 self._init_graph_network(self.inputs, self.outputs, name=self.name)
--> 277 return super(Sequential, self).call(inputs, training=training, mask=mask)
278
279 outputs = inputs # handle the corner case where self.layers is empty
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/network.py in call(self, inputs, training, mask)
715 ' implement a `call` method.')
716
--> 717 return self._run_internal_graph(
718 inputs, training=training, mask=mask,
719 convert_kwargs_to_constants=base_layer_utils.call_context().saving)
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/network.py in _run_internal_graph(self, inputs, training, mask, convert_kwargs_to_constants)
886
887 # Compute outputs.
--> 888 output_tensors = layer(computed_tensors, **kwargs)
889
890 # Update tensor_dict.
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/layers/recurrent.py in __call__(self, inputs, initial_state, constants, **kwargs)
652
653 if initial_state is None and constants is None:
--> 654 return super(RNN, self).__call__(inputs, **kwargs)
655
656 # If any of `initial_state` or `constants` are specified and are Keras
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
966 with base_layer_utils.autocast_context_manager(
967 self._compute_dtype):
--> 968 outputs = self.call(cast_inputs, *args, **kwargs)
969 self._handle_activity_regularization(inputs, outputs)
970 self._set_mask_metadata(inputs, outputs, input_masks)
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/layers/recurrent_v2.py in call(self, inputs, mask, training, initial_state)
1178 # GPU implementation when GPU is available.
1179 if can_use_gpu:
-> 1180 last_output, outputs, new_h, new_c, runtime = gpu_lstm(
1181 **gpu_lstm_kwargs)
1182 else:
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/layers/recurrent_v2.py in gpu_lstm(inputs, init_h, init_c, kernel, recurrent_kernel, bias, mask, time_major, go_backwards, sequence_lengths)
1404 inputs = array_ops.reverse_sequence_v2(
1405 inputs, sequence_lengths, seq_axis=seq_axis, batch_axis=batch_axis)
-> 1406 outputs, h, c, _, _ = gen_cudnn_rnn_ops.cudnn_rnnv3(
1407 inputs,
1408 input_h=init_h,
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/ops/gen_cudnn_rnn_ops.py in cudnn_rnnv3(input, input_h, input_c, params, sequence_lengths, rnn_mode, input_mode, direction, dropout, seed, seed2, num_proj, is_training, time_major, name)
1899 except _core._FallbackException:
1900 try:
-> 1901 return cudnn_rnnv3_eager_fallback(
1902 input, input_h, input_c, params, sequence_lengths,
1903 rnn_mode=rnn_mode, input_mode=input_mode, direction=direction,
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/ops/gen_cudnn_rnn_ops.py in cudnn_rnnv3_eager_fallback(input, input_h, input_c, params, sequence_lengths, rnn_mode, input_mode, direction, dropout, seed, seed2, num_proj, is_training, time_major, name, ctx)
1999 "direction", direction, "dropout", dropout, "seed", seed, "seed2", seed2,
2000 "num_proj", num_proj, "is_training", is_training, "time_major", time_major)
-> 2001 _result = _execute.execute(b"CudnnRNNV3", 5, inputs=_inputs_flat,
2002 attrs=_attrs, ctx=ctx, name=name)
2003 if _execute.must_record_gradient():
~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
57 try:
58 ctx.ensure_initialized()
---> 59 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
60 inputs, attrs, num_outputs)
61 except core._NotOkStatusException as e:
InvalidArgumentError: Invalid input_h shape: [1,8,8] [10,8,8] [Op:CudnnRNNV3]
The 10 refers to the number of units in the LSTM-layer, which is equal to the bounding shape of t1. The two 8's refer to batch_size and window_length. I thought that 1 refered to the output shape, but that is not the case, since it does not change when i add more units to the Dense-layer the number remains the same.

Answer 1

When working with tf.RaggedTensor with sequences of variable sizes, you want to set the batch_size = 1 and you want to ensure that sequence_length passed to the LSTM is None.
This is because even though tf.RaggedTensor are a great storage for variable sized sequences in form of numpy arrays, the LSTM still expects the same length sequence for each batch. You can however have variable sized sequences across batches.
Making those changes should fix the issue you are facing.

Answer 2

This is a bug I encountered as well. It is specific to combination of the below:
Keras implementation of RNN / LSTM Tensorflow version Using GPU And using RaggedTensor
The erroneous source code is comparing shape of ragged tensor of input hidden state of RNN cell with shape of hidden state
The ragged dimension is shown as 1 in the error message.
I was able to replicate same error using your code and Tensorflow 2.2.0
Changing code to remove ragged tensor or running on cpu or different tensorflow version I was able to get this code to work