Метод bucket_by_sequence_length

Метод bucket_by_sequence_length применяется к набору данных Dataset и позволяет эффективно группировать элементы по длине последовательностей. Это особенно полезно при работе с текстовыми или временными рядами, где длины примеров сильно различаются. Первым параметром передается функция element_length_func, которая вычисляет длину каждого элемента. Вторым параметром задается список границ корзин bucket_boundaries. Третьим параметром можно передать размер пакета batch_size. Дополнительно можно указать функцию для формирования ключа сортировки bucket_key и функцию для заполнения пакетов padding_values.

Синтаксис

Dataset.bucket_by_sequence_length(
	element_length_func,
	bucket_boundaries,
	bucket_batch_sizes,
	padding_values=None,
	pad_to_bucket_boundary=False,
	no_padding=False,
	drop_remainder=False
)

Пример

Давайте создадим набор данных из последовательностей разной длины и сгруппируем их по корзинам:

import tensorflow as tf

tf.random.set_seed(0)

data = [
	[1, 2, 3],
	[4, 5],
	[6, 7, 8, 9],
	[10],
	[11, 12, 13, 14, 15]
]

dataset = tf.data.Dataset.from_tensor_slices(data)

batched = dataset.bucket_by_sequence_length(
	element_length_func=lambda x: tf.shape(x)[0],
	bucket_boundaries=[2, 4],
	bucket_batch_sizes=[2, 2, 2],
	padding_values=tf.constant(0, dtype=tf.int32),
	pad_to_bucket_boundary=True
)

for batch in batched:
	print(batch)

Результат выполнения кода:

tf.Tensor(
[[ 4  5  0  0]
 [10  0  0  0]], shape=(2, 4), dtype=int32)
tf.Tensor(
[[ 1  2  3  0]
 [ 6  7  8  9]], shape=(2, 4), dtype=int32)
tf.Tensor(
[[11 12 13 14 15]], shape=(1, 5), dtype=int32)

Пример

Давайте рассмотрим пример без заполнения границ корзин, с отбрасыванием неполных пакетов:

import tensorflow as tf

tf.random.set_seed(0)

data = [
	[1, 2, 3],
	[4, 5],
	[6, 7, 8, 9],
	[10],
	[11, 12, 13, 14, 15]
]

dataset = tf.data.Dataset.from_tensor_slices(data)

batched = dataset.bucket_by_sequence_length(
	element_length_func=lambda x: tf.shape(x)[0],
	bucket_boundaries=[2, 4],
	bucket_batch_sizes=[1, 1, 1],
	padding_values=tf.constant(0, dtype=tf.int32),
	drop_remainder=True
)

for batch in batched:
	print(batch)

Результат выполнения кода:

tf.Tensor([[4 5]], shape=(1, 2), dtype=int32)
tf.Tensor([[1 2 3]], shape=(1, 3), dtype=int32)
tf.Tensor([[6 7 8 9]], shape=(1, 4), dtype=int32)
tf.Tensor([[10]], shape=(1, 1), dtype=int32)

Смотрите также

класс Dataset,
который представляет набор данных
метод padded_batch,
который формирует пакеты с заполнением до одинаковой длины
метод group_by_window,
который группирует элементы по окнам
метод batch,
который объединяет элементы в пакеты