Python tf.data.Dataset.padded_batch用法及代码示例

用法

padded_batch(
    batch_size, padded_shapes=None, padding_values=None, drop_remainder=False,
    name=None
)

将此数据集的连续元素组合成填充批次。

此转换将输入数据集的多个连续元素组合成一个元素。

与 tf.data.Dataset.batch 一样，结果元素的组件将有一个额外的外部维度，即 batch_size (如果 batch_size 不均匀地划分输入元素的数量 N ，则为最后一个元素的 N % batch_size 并且drop_remainder 是 False )。如果您的程序依赖于具有相同外部尺寸的批次，则应将 drop_remainder 参数设置为 True 以防止生成较小的批次。

与 tf.data.Dataset.batch 不同，要批处理的输入元素可能具有不同的形状，并且此转换会将每个组件填充到 padded_shapes 中的相应形状。 padded_shapes 参数确定输出元素中每个组件的每个维度的结果形状：

• 如果尺寸是常数，则组件将在该尺寸中填充到该长度。
• 如果维度未知，组件将被填充到该维度中所有元素的最大长度。

A = (tf.data.Dataset
     .range(1, 5, output_type=tf.int32)
     .map(lambda x:tf.fill([x], x)))
# Pad to the smallest per-batch size that fits all elements.
B = A.padded_batch(2)
for element in B.as_numpy_iterator():
  print(element)
[[1 0]
 [2 2]]
[[3 3 3 0]
 [4 4 4 4]]
# Pad to a fixed size.
C = A.padded_batch(2, padded_shapes=5)
for element in C.as_numpy_iterator():
  print(element)
[[1 0 0 0 0]
 [2 2 0 0 0]]
[[3 3 3 0 0]
 [4 4 4 4 0]]
# Pad with a custom value.
D = A.padded_batch(2, padded_shapes=5, padding_values=-1)
for element in D.as_numpy_iterator():
  print(element)
[[ 1 -1 -1 -1 -1]
 [ 2  2 -1 -1 -1]]
[[ 3  3  3 -1 -1]
 [ 4  4  4  4 -1]]
# Components of nested elements can be padded independently.
elements = [([1, 2, 3], [10]),
            ([4, 5], [11, 12])]
dataset = tf.data.Dataset.from_generator(
    lambda:iter(elements), (tf.int32, tf.int32))
# Pad the first component of the tuple to length 4, and the second
# component to the smallest size that fits.
dataset = dataset.padded_batch(2,
    padded_shapes=([4], [None]),
    padding_values=(-1, 100))
list(dataset.as_numpy_iterator())
[(array([[ 1,  2,  3, -1], [ 4,  5, -1, -1]], dtype=int32),
  array([[ 10, 100], [ 11,  12]], dtype=int32))]
# Pad with a single value and multiple components.
E = tf.data.Dataset.zip((A, A)).padded_batch(2, padding_values=-1)
for element in E.as_numpy_iterator():
  print(element)
(array([[ 1, -1],
       [ 2,  2]], dtype=int32), array([[ 1, -1],
       [ 2,  2]], dtype=int32))
(array([[ 3,  3,  3, -1],
       [ 4,  4,  4,  4]], dtype=int32), array([[ 3,  3,  3, -1],
       [ 4,  4,  4,  4]], dtype=int32))

另请参见 tf.data.experimental.dense_to_sparse_batch ，它将可能具有不同形状的元素组合成 tf.sparse.SparseTensor 。