Tensorflow - Pad OR Truncate Sequence with Dataset API

user_2399935 0
Viewed 2172

I am trying to use the Dataset API to prepare a TFRecordDataset of text sequences. After processing, I have a dictionary of tensors for each record. Each record contains two sequences.

I am using padded_batch to apply padding

dataset = dataset.padded_batch(batch_size, padded_shapes= {
    'seq1': tf.TensorShape([None]),
    'seq2': tf.TensorShape([None])
})

This pads each sequence to the maximum sequence lengths in the batch. However, I would like to choose an arbitrary sequence length and pad to this length when the true sequence length is less else truncate the sequence.

When I try replacing None with 100 for example, I encounter a DataLossError

DataLossError: Attempted to pad to a smaller size than the input element.

Is there a way to do this to achieve similar functionality as tf.image.resize_image_with_crop_or_pad on a sequence?

tensorflow-datasets tensorflow sequence
2 Answers

There's no easy way to pad or truncate, but you can use the map function to obtain a dataset containing elements with the desired length. Here's a quick example:

k = 4
def pad_or_trunc(t):
    dim = tf.size(t)
    return tf.cond(tf.equal(dim, k), lambda: t, lambda: tf.cond(tf.greater(dim, k), lambda: tf.slice(t, [0], [k]), lambda: tf.concat([t, tf.zeros(k-dim, dtype=tf.int32)], 0)))

vals = tf.constant([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
dset1 = tf.data.Dataset.from_tensor_slices(vals)
dset2 = dset1.map(pad_or_trunc)
iter = dset2.make_one_shot_iterator()

with tf.Session() as sess:
    while True:
        try:
            print(sess.run(iter.get_next()))
        except tf.errors.OutOfRangeError:
            break
user_9349040

You can first truncate all longer sequences using tf.slice and tf.math.greater, and then use padded_batch to pad the sequences.

An example could look like this:

import tensorflow as tf
import numpy as np

# data generator
def gen():
  for i in [np.array([1, 1, 1]), np.array([2, 2, 2, 2]), np.array([3, 3, 3, 3, 3])]:
    yield i

cut_or_pad = 4 # 100 in your example

def cut_if_longer(el):
  if tf.greater(tf.shape(el), cut_or_pad): # only slice if longer
    return tf.slice(el, begin=[0], size=[cut_or_pad])
  return el

# data pipeline
dataset = tf.data.Dataset.from_generator( gen, (tf.int32), (tf.TensorShape([None])))
dataset = dataset.map( lambda el: cut_if_longer(el))
dataset = dataset.padded_batch(batch_size=2, padded_shapes=[cut_or_pad], padding_values=-1)

list(dataset.take(2).as_numpy_iterator())
user_9510430
Related