How to create buckets using tf.data.experimental.bucket_by_sequence_length, and each bucket should be padded to the length of its coresponding boundry

user_14210738 0
Viewed 102

I would like to know how to make data bucketting using tf.data.experimental.bucket_by_sequence_length, but I want my outputs to be of shapes (bucket_batch_sizes[i], bucket_boundaries[i]) Here is my code

import tensorflow as tf
def elements_gen():
   text = [[1, 2, 3], [3, 4, 5, 6, 7], [1, 2], [8, 9, 0, 2]]
   label = [1, 2, 1, 2]
   for x, y in zip(text, label):
       yield (x, y)

def element_length_fn(x, y):
   return tf.shape(x)[0]

dataset = tf.data.Dataset.from_generator(generator=elements_gen,
                                     output_shapes=([None],[]),
                                     output_types=(tf.int32, tf.int32))

dataset =   dataset.apply(tf.data.experimental.bucket_by_sequence_length(element_length_func=element_length_fn, padding_values=0,bucket_batch_sizes=[2, 2, 2,2],bucket_boundaries=[0, 3, 8]))

for elem in dataset.as_numpy_iterator():
  print(elem[0].shape)
  print(elem)

this is the output

(2, 5)
(array([[1, 2, 3, 0, 0],
       [3, 4, 5, 6, 7]], dtype=int32), array([1, 2], dtype=int32))
(1, 2)
(array([[1, 2]], dtype=int32), array([1], dtype=int32))
(1, 4)
(array([[8, 9, 0, 2]], dtype=int32), array([2], dtype=int32))
tensorflow bucket python
1 Answers

the issue was solved by setting the attribute pad_to_bucket_boundary=True in tf.data.experimental.bucket_by_sequence_length as follow

dataset =   dataset.apply(tf.data.experimental.bucket_by_sequence_length(element_length_func=element_length_fn,
                                                              padding_values=0,
                                                              bucket_batch_sizes=[2, 2, 2,2],
                                                              pad_to_bucket_boundary=True,
                                                              bucket_boundaries=[0, 3, 8]))
user_142107380
Related