添加具有未知大小张量的 TensorFlow 数据集

Question

我想向 TFDS 添加一个 TensorFlow 数据集，其中包含一堆形状未知的 2 阶张量。要定义功能，我使用以下功能tfds.features.Tensor(shape=[None, None], dtype=tf.float32)。但是，当我尝试加载数据集时，它失败并出现错误：

NotImplementedError: Specification error for feature k (TensorInfo(shape=[None, None], dtype=tf.float32)): Tensor with a unknown dimension not at the first position not supported: TensorInfo(shape=[None, None], dtype=tf.float32)

重现错误的最小示例 (Real example on Github)：

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds


class RandomTensors(tfds.core.GeneratorBasedBuilder):

    VERSION = tfds.core.Version('0.1.0')

    def _info(self):
        return tfds.core.DatasetInfo(
            builder=self,
            description="",
            features=tfds.features.FeaturesDict({
                'k':
                tfds.features.Tensor(shape=[None, None], dtype=tf.float32)
            }),
            homepage="",
            citation="",
        )

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        return [
            tfds.core.SplitGenerator(
                name=tfds.Split.TEST,
                gen_kwargs={},
            ),
        ]

    def _generate_examples(self):
        for idx in range(100):
            size = np.random.uniform(10, 20, size=(2, )).astype(np.int32)
            k = np.random.normal(size=size).astype(np.float32)
            yield idx, {'k': k}


if __name__ == '__main__':
    d = tfds.load('random_tensors', split='test')

TensorFlow 数据集中是否不支持维度未知的张量？还是我的代码有错误？

Answer 1

我认为这不可能与 TensorFlow 数据集有关，因为正如错误消息所解释的那样，只有第一个维度（通常是批处理维度）可以是动态的。

与此相关，tf.data.Dataset 对象通常需要一个矩形数组。以下失败，例如：

tf.data.Dataset.from_tensor_slices([[1], [1, 2]])

这可以直接使用tf.data.Dataset 对象进行，但您将无法以 >= 2 的批次大小对样本进行批处理，因为它们的形状不同。

def _generate_examples():
    for idx in range(100):
        size = np.random.uniform(10, 20, size=(2, )).astype(np.int32)
        k = np.random.normal(size=size).astype(np.float32)
        yield k

d = tf.data.Dataset.from_generator(_generate_examples, output_types=np.float32)
next(iter(d))  # Works!

next(iter(d.batch(2)))
# Fails with error: 
# InvalidArgumentError: Cannot batch tensors with different shapes in component 0.
# First element had shape [11,10] and element 1 had shape [13,17].

---

您可能想探索使用具有静态形状的较大张量，然后用动态大小的较小张量填充它的可能性。这样你就有了一个静态的形状，但仍然有动态形状的数据。动态数据的形状可以是数据集的一个特征，它与索引和数据一起在生成器中生成。