如何在 Tensorflow 2.0 中制作参差不齐的批次?

How do I make a ragged batch in Tensorflow 2.0?

我正在尝试从由一维数值数据张量组成的 Tensorflow 数据集创建数据输入管道。我想创建一批参差不齐的张量;我不想填充数据。

例如,如果我的数据是以下形式:

[
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 
    [0, 1, 2, 3, 4]
    ...
]

我希望我的数据集包含以下形式的批次:

<tf.Tensor [
    <tf.RaggedTensor [
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 
        [0, 1, 2, 3, 4], 
        ...]>,
    <tf.RaggedTensor [
        [ ... ],
        ...]>
    ]>

我试过使用地图创建 RaggedTensor,但我似乎无法在一维数据上创建它。

我认为这可以通过批处理前后的一些工作来实现。

# First, you can expand along the 0 axis for each data point
dataset = dataset.map(lambda x: tf.expand_dims(x, 0))
# Then create a RaggedTensor with a ragged rank of 1
dataset = dataset.map(lambda x: tf.RaggedTensor.from_tensor(x))
# Create batches
dataset = dataset.batch(BATCH_SIZE)
# Squeeze the extra dimension from the created batches
dataset = dataset.map(lambda x: tf.squeeze(x, axis=1))

那么最终的输出将是这样的形式:

<tf.RaggedTensor [
    <tf.Tensor [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]>,
    <tf.Tensor [0, 1, 2, 3]>,
    ...
]>

每批。