Python 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
)

警告：此函数已弃用。它将在未来的版本中删除。更新说明：使用 tf.data.Dataset.bucket_by_sequence_length(...) 。

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]]