我有一个包含 2 阶张量的 Tensorflow 数据集。每个张量都是从文件中读取的，具有变化的第一维和恒定的第二维。例如，每个项目的形状可以是以下形式：

(123, 3)
(100, 3)
(134, 3)
...

为了模拟类似的行为，您可以使用以下命令：

lengths = tf.random.uniform(shape=(10,), minval=5, maxval=10, dtype=tf.int32)
# this would be a dataset of filenames in my code
dataset = tf.data.Dataset.from_tensor_slices(lengths)

def fake_read_file(tensor):
    # this function would read and preprocess data from file in my code
    # however, for demonstration purposes it just returns a dummy data with varying length in the first dimension
    dummy_data = tf.convert_to_tensor([[0.2, 0.2, 0.2]])
    return tf.repeat(dummy_data, tensor, axis=0)

# I use tf.py_function decorator because file read can be done only in eager mode
dataset = dataset.map(lambda x: tf.py_function(fake_read_file, inp=[x], Tout=tf.float32))

我需要使用 RaggedTensors 批处理此数据集（我不想填充输入）。每批的尺寸应为(batch_size, None, 3)。

现在我想到了以下解决方案：

dataset = dataset.map(lambda x: tf.RaggedTensor.from_tensor(x))
dataset = dataset.batch(batch_size=2, drop_remainder=True)

然而，当我迭代数据集时，我得到了形状的批次(batch_size, None, None)，即最后一个维度也变得参差不齐。

解决方案

我发现解决这个问题的方法是在将张量转换为不规则张量之前添加一个虚拟维度，然后在批处理之后删除该维度：

dataset = dataset.map(lambda x: tf.expand_dims(x, 0))
dataset = dataset.map(lambda x: tf.RaggedTensor.from_tensor(x))
dataset = dataset.batch(batch_size=2, drop_remainder=True)
dataset = dataset.map(lambda x: tf.squeeze(x, axis=1))

我不知道 Tensorflow 为何会这样，但现在它能生成正确的批次。此外，当您迭代批次中的元素时，您将获得常规张量（形状不同），而不是像以前那样不规则的张量。