tf.data.experimental.bucket_by_sequence_length

Преобразование, которое группирует элементы в Dataset по длине. (устаревшее)

tf.data.experimental.bucket_by_sequence_length(
    element_length_func,
    bucket_boundaries,
    bucket_batch_sizes,
    padded_shapes=None,
    padding_values=None,
    pad_to_bucket_boundary=False,
    no_padding=False,
    drop_remainder=False
)

Элементы Dataset группируются по длине, а затем заполняются и объединяются в пакеты.

Это полезно для задач с последовательностями, в которых элементы имеют переменную длину. Группирование элементов с похожими длинами уменьшает общую долю заполнения в пакете, что повышает эффективность шагов обучения.

Ниже приведен пример букетизации входных данных в 3 корзины "[0, 3), [3, 5), [5, inf)" на основе длины последовательности с размером пакета 2.

elements = [
  [0], [1, 2, 3, 4], [5, 6, 7],
  [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]

dataset = tf.data.Dataset.from_generator(
    lambda: elements, tf.int64, output_shapes=[None])

dataset = dataset.apply(
    tf.data.experimental.bucket_by_sequence_length(
        element_length_func=lambda elem: tf.shape(elem)[0],
        bucket_boundaries=[3, 5],
        bucket_batch_sizes=[2, 2, 2]))

for elem in dataset.as_numpy_iterator():
  print(elem)
[[1 2 3 4]
 [5 6 7 0]]
[[ 7  8  9 10 11  0]
 [13 14 15 16 19 20]]
[[ 0  0]
 [21 22]]

Также существует возможность заполнить набор данных до границы корзины. Вы также можете указать значение, которое будет использоваться при заполнении данных. В примере ниже используется -1 в качестве заполнения, а также показано, как входные данные букетизируются в две корзины "[0,3], [4,6]".

elements = [
  [0], [1, 2, 3, 4], [5, 6, 7],
  [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]

dataset = tf.data.Dataset.from_generator(
  lambda: elements, tf.int32, output_shapes=[None])

dataset = dataset.apply(
    tf.data.experimental.bucket_by_sequence_length(
        element_length_func=lambda elem: tf.shape(elem)[0],
        bucket_boundaries=[4, 7],
        bucket_batch_sizes=[2, 2, 2],
        pad_to_bucket_boundary=True,
        padding_values=-1))

for elem in dataset.as_numpy_iterator():
  print(elem)
[[ 0 -1 -1]
 [ 5  6  7]]
[[ 1  2  3  4 -1 -1]
 [ 7  8  9 10 11 -1]]
[[21 22 -1]]
[[13 14 15 16 19 20]]

При использовании опции pad_to_bucket_boundary, может оказаться, что поддерживать размер пакета корзины не всегда возможно. Вы можете отбросить пакеты, которые не поддерживают размер пакета корзины, используя опцию drop_remainder. Используя те же входные данные, что и в примере выше, вы получите следующий результат.

elements = [
  [0], [1, 2, 3, 4], [5, 6, 7],
  [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]

dataset = tf.data.Dataset.from_generator(
  lambda: elements, tf.int32, output_shapes=[None])

dataset = dataset.apply(
    tf.data.experimental.bucket_by_sequence_length(
        element_length_func=lambda elem: tf.shape(elem)[0],
        bucket_boundaries=[4, 7],
        bucket_batch_sizes=[2, 2, 2],
        pad_to_bucket_boundary=True,
        padding_values=-1,
        drop_remainder=True))

for elem in dataset.as_numpy_iterator():
  print(elem)
[[ 0 -1 -1]
 [ 5  6  7]]
[[ 1  2  3  4 -1 -1]
 [ 7  8  9 10 11 -1]]