Загрузите внешний tfrecord с помощью TFDS

Если у вас есть прототип tf.train.Example (внутри .tfrecord , .riegeli ,...), созданный сторонними инструментами и который вы хотели бы загрузить напрямую с помощью tfds API, то эта страница для вас.

Чтобы загрузить файлы .tfrecord , вам нужно всего лишь:

  • Следуйте соглашению об именах TFDS.
  • Добавьте файлы метаданных ( dataset_info.json , features.json ) в файлы tfrecord.

Ограничения:

Соглашение об именах файлов

TFDS поддерживает определение шаблона имен файлов, что обеспечивает гибкость использования различных схем именования файлов. Шаблон представлен tfds.core.ShardedFileTemplate и поддерживает следующие переменные: {DATASET} , {SPLIT} , {FILEFORMAT} , {SHARD_INDEX} , {NUM_SHARDS} и {SHARD_X_OF_Y} . Например, схема именования файлов по умолчанию в TFDS: {DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y} . Для MNIST это означает, что имена файлов выглядят следующим образом:

  • mnist-test.tfrecord-00000-of-00001
  • mnist-train.tfrecord-00000-of-00001

Добавить метаданные

Предоставьте структуру функций

Чтобы TFDS мог декодировать прототип tf.train.Example , вам необходимо предоставить структуру tfds.features соответствующую вашим спецификациям. Например:

features = tfds.features.FeaturesDict({
    'image':
        tfds.features.Image(
            shape=(256, 256, 3),
            doc='Picture taken by smartphone, downscaled.'),
    'label':
        tfds.features.ClassLabel(names=['dog', 'cat']),
    'objects':
        tfds.features.Sequence({
            'camera/K': tfds.features.Tensor(shape=(3,), dtype=tf.float32),
        }),
})

Соответствует следующим спецификациям tf.train.Example :

{
    'image': tf.io.FixedLenFeature(shape=(), dtype=tf.string),
    'label': tf.io.FixedLenFeature(shape=(), dtype=tf.int64),
    'objects/camera/K': tf.io.FixedLenSequenceFeature(shape=(3,), dtype=tf.int64),
}

Указание функций позволяет TFDS автоматически декодировать изображения, видео и т. д. Как и любые другие наборы данных TFDS, метаданные функций (например, имена меток,...) будут доступны пользователю (например, info.features['label'].names ).

Если вы контролируете конвейер генерации

Если вы генерируете наборы данных вне TFDS, но по-прежнему контролируете конвейер генерации, вы можете использовать tfds.features.FeatureConnector.serialize_example для кодирования ваших данных из dict[np.ndarray] в tf.train.Example proto bytes :

with tf.io.TFRecordWriter('path/to/file.tfrecord') as writer:
  for ex in all_exs:
    ex_bytes = features.serialize_example(data)
    writer.write(ex_bytes)

Это обеспечит совместимость функций с TFDS.

Аналогично, существует feature.deserialize_example для декодирования прототипа ( пример ).

Если вы не контролируете конвейер генерации

Если вы хотите увидеть, как tfds.features представлены в tf.train.Example , вы можете изучить это в colab:

  • Чтобы перевести tfds.features в удобочитаемую структуру tf.train.Example , вы можете вызвать features.get_serialized_info() .
  • Чтобы получить точную спецификацию FixedLenFeature ,..., переданную в tf.io.parse_single_example , вы можете использовать spec = features.tf_example_spec

Получить статистику по сплитам

TFDS требует знать точное количество примеров в каждом сегменте. Это необходимо для таких функций, как len(ds) или API subplit : split='train[75%:]' .

  • Если у вас есть эта информация, вы можете явно создать список tfds.core.SplitInfo и перейти к следующему разделу:

    split_infos = [
        tfds.core.SplitInfo(
            name='train',
            shard_lengths=[1024, ...],  # Num of examples in shard0, shard1,...
            num_bytes=0,  # Total size of your dataset (if unknown, set to 0)
        ),
        tfds.core.SplitInfo(name='test', ...),
    ]
    
  • Если вы не знаете эту информацию, вы можете вычислить ее с помощью сценария compute_split_info.py (или вашего собственного сценария с помощью tfds.folder_dataset.compute_split_info ). Он запустит лучевой конвейер, который будет читать все сегменты в заданном каталоге и вычислять информацию.

Добавить файлы метаданных

Чтобы автоматически добавлять нужные файлы метаданных в ваш набор данных, используйте tfds.folder_dataset.write_metadata :

tfds.folder_dataset.write_metadata(
    data_dir='/path/to/my/dataset/1.0.0/',
    features=features,
    # Pass the `out_dir` argument of compute_split_info (see section above)
    # You can also explicitly pass a list of `tfds.core.SplitInfo`.
    split_infos='/path/to/my/dataset/1.0.0/',
    # Pass a custom file name template or use None for the default TFDS
    # file name template.
    filename_template='{SPLIT}-{SHARD_X_OF_Y}.{FILEFORMAT}',

    # Optionally, additional DatasetInfo metadata can be provided
    # See:
    # https://www.tensorflow.org/datasets/api_docs/python/tfds/core/DatasetInfo
    description="""Multi-line description."""
    homepage='http://my-project.org',
    supervised_keys=('image', 'label'),
    citation="""BibTex citation.""",
)

После того как функция была вызвана один раз в каталоге вашего набора данных, файлы метаданных ( dataset_info.json ,...) были добавлены, и ваши наборы данных готовы к загрузке с помощью TFDS (см. следующий раздел).

Загрузить набор данных с помощью TFDS

Прямо из папки

После создания метаданных наборы данных можно загрузить с помощью tfds.builder_from_directory , который возвращает tfds.core.DatasetBuilder со стандартным API TFDS (например, tfds.builder ):

builder = tfds.builder_from_directory('~/path/to/my_dataset/3.0.0/')

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...

Непосредственно из нескольких папок

Также возможно загружать данные из нескольких папок. Это может произойти, например, при обучении с подкреплением, когда каждый из нескольких агентов генерирует отдельный набор данных, и вы хотите загрузить их все вместе. Другие варианты использования — это когда новый набор данных создается на регулярной основе, например, новый набор данных каждый день, и вы хотите загрузить данные из диапазона дат.

Чтобы загрузить данные из нескольких папок, используйте tfds.builder_from_directories , который возвращает tfds.core.DatasetBuilder со стандартным API TFDS (например, tfds.builder ):

builder = tfds.builder_from_directories(builder_dirs=[
    '~/path/my_dataset/agent1/1.0.0/',
    '~/path/my_dataset/agent2/1.0.0/',
    '~/path/my_dataset/agent3/1.0.0/',
])

# Metadata are available as usual
builder.info.splits['train'].num_examples

# Construct the tf.data.Dataset pipeline
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
  ...

Структура папок (необязательно)

Для лучшей совместимости с TFDS вы можете организовать свои данные как <data_dir>/<dataset_name>[/<dataset_config>]/<dataset_version> . Например:

data_dir/
    dataset0/
        1.0.0/
        1.0.1/
    dataset1/
        config0/
            2.0.0/
        config1/
            2.0.0/

Это сделает ваши наборы данных совместимыми с API tfds.load / tfds.builder , просто указав data_dir/ :

ds0 = tfds.load('dataset0', data_dir='data_dir/')
ds1 = tfds.load('dataset1/config0', data_dir='data_dir/')