how to pad sequences in a tensor slice dataset in TensorFlow?

Question

I have a tensor slice dataset made from two ragged tensors.

tensor_a is like: <tf.RaggedTensor [[3, 3, 5], [3, 3, 14, 4, 17, 20], [3, 14, 22, 17]]>

tensor_b is like: <tf.RaggedTensor [[-1, 1, -1], [-1, -1, 1, -1, -1, -1], [-1, 1, -1, 2]]>

(Same index, same length for tensor_a and tensor_b.)

I made the dataset by

dataset = tf.data.Dataset.from_tensor_slices((tensor_a, tensor_b))
dataset

<TensorSliceDataset element_spec=(RaggedTensorSpec(TensorShape([None]), tf.int64, 0, tf.int64), RaggedTensorSpec(TensorShape([None]), tf.int32, 0, tf.int64))>

How to pad the sequences in my dataset? I've tried tf.pad and tf.keras.preprocessing.sequence.pad_sequences but haven't found a right way.

Answer 1

You could try something like this:

import tensorflow as tf

tensor_a = tf.ragged.constant([[3, 3, 5], [3, 3, 14, 4, 17, 20], [3, 14, 22, 17]])
tensor_b = tf.ragged.constant([[-1, 1, -1], [-1, -1, 1, -1, -1, -1], [-1, 1, -1, 2]])
dataset = tf.data.Dataset.from_tensor_slices((tensor_a, tensor_b))

max_length = max(list(dataset.map(lambda x, y: tf.shape(x)[0])))

def pad(x, y):
  x = tf.concat([x, tf.zeros((int(max_length-tf.shape(x)[0]),), dtype=tf.int32)], axis=0)
  y = tf.concat([y, tf.zeros((int(max_length-tf.shape(y)[0]),), dtype=tf.int32)], axis=0)
  return x, y

dataset = dataset.map(pad)
for x, y in dataset:
  print(x, y)

tf.Tensor([3 3 5 0 0 0], shape=(6,), dtype=int32) tf.Tensor([-1  1 -1  0  0  0], shape=(6,), dtype=int32)
tf.Tensor([ 3  3 14  4 17 20], shape=(6,), dtype=int32) tf.Tensor([-1 -1  1 -1 -1 -1], shape=(6,), dtype=int32)
tf.Tensor([ 3 14 22 17  0  0], shape=(6,), dtype=int32) tf.Tensor([-1  1 -1  2  0  0], shape=(6,), dtype=int32)

For pre-padding, just adjust the pad function:

def pad(x, y):
  x = tf.concat([tf.zeros((int(max_length-tf.shape(x)[0]),), dtype=tf.int32), x], axis=0)
  y = tf.concat([tf.zeros((int(max_length-tf.shape(y)[0]),), dtype=tf.int32), y], axis=0)
  return x, y

tf.Tensor([0 0 0 3 3 5], shape=(6,), dtype=int32) tf.Tensor([ 0  0  0 -1  1 -1], shape=(6,), dtype=int32)
tf.Tensor([ 3  3 14  4 17 20], shape=(6,), dtype=int32) tf.Tensor([-1 -1  1 -1 -1 -1], shape=(6,), dtype=int32)
tf.Tensor([ 0  0  3 14 22 17], shape=(6,), dtype=int32) tf.Tensor([ 0  0 -1  1 -1  2], shape=(6,), dtype=int32)