tf.data.experimental.dense_to_ragged_batch的功能

本文根据官网手册所写，简而言之该函数就是将数据集中的每batch_size个元素组成一个tf.RaggedTensor类型的多个数据。从下面的例程中可以比较直观的看出来

import tensorflow as tf
import numpy as np

dataset = tf.data.Dataset.from_tensor_slices(np.arange(6))
for batch in dataset:
  print(batch)

打印

tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(2, shape=(), dtype=int64)
tf.Tensor(3, shape=(), dtype=int64)
tf.Tensor(4, shape=(), dtype=int64)
tf.Tensor(5, shape=(), dtype=int64)

dataset = dataset.map(lambda x: tf.range(x))
for batch in dataset:
  print(batch)

打印

tf.Tensor([], shape=(0,), dtype=int64)
tf.Tensor([0], shape=(1,), dtype=int64)
tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor([0 1 2], shape=(3,), dtype=int64)
tf.Tensor([0 1 2 3], shape=(4,), dtype=int64)
tf.Tensor([0 1 2 3 4], shape=(5,), dtype=int64)

dataset.element_spec.shape
for batch in dataset:
  print(batch)

打印

tf.Tensor([], shape=(0,), dtype=int64)
tf.Tensor([0], shape=(1,), dtype=int64)
tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor([0 1 2], shape=(3,), dtype=int64)
tf.Tensor([0 1 2 3], shape=(4,), dtype=int64)
tf.Tensor([0 1 2 3 4], shape=(5,), dtype=int64)

当batch_size=5，drop_remainder=True时，输出如下

dataset = dataset.apply(
    tf.data.experimental.dense_to_ragged_batch(batch_size=5, drop_remainder=True))
for batch in dataset:
  print(batch)

打印如下，可以看到将五个元素组合成一个tf.RaggedTensor数据，且把后面不足5个的部分丢弃了。

<tf.RaggedTensor [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]]>

当batch_size=3，drop_remainder=True时，输出如下

dataset = dataset.apply(
    tf.data.experimental.dense_to_ragged_batch(batch_size=3, drop_remainder=True))
for batch in dataset:
  print(batch)

打印如下，可以看到将3个元素组合成一个tf.RaggedTensor数据，因为没有不足5个的部分也就什么也没丢弃。

<tf.RaggedTensor [[], [0], [0, 1]]>
<tf.RaggedTensor [[0, 1, 2], [0, 1, 2, 3], [0, 1, 2, 3, 4]]>