Add a TensorFlow Dataset with Tensors of unknown size

user_12973831 0
Viewed 847

I would like to add a TensorFlow dataset to TFDS which contains a bunch of rank 2 tensors with unknown shapes. To define the features I use the following feature tfds.features.Tensor(shape=[None, None], dtype=tf.float32). However, when I try to load the dataset it fails with the error:

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)

A minimal example to reproduce the error (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')

Are tensors with unknown dimensions not supported in TensorFlow Datasets? Or do I have an an error in my code?

tensorflow-datasets tensorflow python
1 Answers

I don't think this is possible to do with TensorFlow Datasets, because, as the error message explains, only the first dimension (typically the batch dimension) can be dynamic.

Related to this, the tf.data.Dataset object typically expects a rectangular array. The following fails, for example:

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

This would be possible to do directly with a tf.data.Dataset object, but you would not be able to batch the samples in a batch size >= 2 because their shapes are different.

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].

You might want to explore the possibility of using a larger tensor with a static shape, and then filling that with a dynamically sized smaller tensor. That way you have a static shape, but still have dynamically shaped data. The shape of the dynamic data can be a feature of the dataset, which is yielded in the generator along with the index and data.

user_56660870
Related