Carregar tfrecord externo com TFDS

Se você possui um proto tf.train.Example (dentro de .tfrecord , .riegeli ,...), que foi gerado por ferramentas de terceiros, que você gostaria de carregar diretamente com a API tfds, então esta página é para você.

Para carregar seus arquivos .tfrecord , você só precisa:

  • Siga a convenção de nomenclatura TFDS.
  • Adicione arquivos de metadados ( dataset_info.json , features.json ) junto com seus arquivos tfrecord.

Limitações:

Convenção de nomenclatura de arquivos

TFDS suporta a definição de um modelo para nomes de arquivos, o que fornece flexibilidade para usar diferentes esquemas de nomenclatura de arquivos. O modelo é representado por um tfds.core.ShardedFileTemplate e suporta as seguintes variáveis: {DATASET} , {SPLIT} , {FILEFORMAT} , {SHARD_INDEX} , {NUM_SHARDS} e {SHARD_X_OF_Y} . Por exemplo, o esquema de nomenclatura de arquivo padrão do TFDS é: {DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y} . Para MNIST, isso significa que os nomes dos arquivos têm a seguinte aparência:

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

Adicionar metadados

Forneça a estrutura de recursos

Para que o TFDS seja capaz de decodificar o proto tf.train.Example , você precisa fornecer a estrutura tfds.features que corresponda às suas especificações. Por exemplo:

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),
        }),
})

Corresponde às seguintes especificações 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),
}

Especificar os recursos permite que o TFDS decodifique automaticamente imagens, vídeos,... Como qualquer outro conjunto de dados TFDS, os metadados dos recursos (por exemplo, nomes de rótulos,...) serão expostos ao usuário (por exemplo, info.features['label'].names ).

Se você controlar o pipeline de geração

Se você gerar conjuntos de dados fora do TFDS, mas ainda controlar o pipeline de geração, poderá usar tfds.features.FeatureConnector.serialize_example para codificar seus dados de dict[np.ndarray] para 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)

Isso garantirá a compatibilidade dos recursos com o TFDS.

Da mesma forma, existe um feature.deserialize_example para decodificar o proto ( exemplo )

Se você não controlar o pipeline de geração

Se você quiser ver como tfds.features são representados em um tf.train.Example , você pode examinar isso no colab:

  • Para traduzir tfds.features na estrutura legível por humanos do tf.train.Example , você pode chamar features.get_serialized_info() .
  • Para obter a especificação exata FixedLenFeature ,... passada para tf.io.parse_single_example , você pode usar spec = features.tf_example_spec

Obtenha estatísticas sobre divisões

O TFDS exige saber o número exato de exemplos dentro de cada fragmento. Isso é necessário para recursos como len(ds) ou a API subplit : split='train[75%:]' .

  • Se você tiver essas informações, poderá criar explicitamente uma lista de tfds.core.SplitInfo e pular para a próxima seção:

    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', ...),
]

  • Se você não souber essas informações, poderá calculá-las usando o script compute_split_info.py (ou em seu próprio script com tfds.folder_dataset.compute_split_info ). Ele iniciará um pipeline de feixe que lerá todos os fragmentos em um determinado diretório e calculará as informações.

Adicionar arquivos de metadados

Para adicionar automaticamente os arquivos de metadados adequados ao seu conjunto de dados, use 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.""",
)

Depois que a função for chamada uma vez em seu diretório de conjunto de dados, os arquivos de metadados ( dataset_info.json ,...) foram adicionados e seus conjuntos de dados estão prontos para serem carregados com TFDS (veja a próxima seção).

Carregar conjunto de dados com TFDS

Diretamente da pasta

Depois que os metadados forem gerados, os conjuntos de dados podem ser carregados usando tfds.builder_from_directory que retorna um tfds.core.DatasetBuilder com a API TFDS padrão (como 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:
  ...

Diretamente de várias pastas

Também é possível carregar dados de várias pastas. Isso pode acontecer, por exemplo, no aprendizado por reforço, quando vários agentes estão gerando, cada um, um conjunto de dados separado e você deseja carregar todos eles juntos. Outros casos de uso são quando um novo conjunto de dados é produzido regularmente, por exemplo, um novo conjunto de dados por dia, e você deseja carregar dados de um intervalo de datas.

Para carregar dados de várias pastas, use tfds.builder_from_directories , que retorna um tfds.core.DatasetBuilder com a API TFDS padrão (como 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:
  ...

Estrutura de pastas (opcional)

Para melhor compatibilidade com TFDS, você pode organizar seus dados como <data_dir>/<dataset_name>[/<dataset_config>]/<dataset_version> . Por exemplo:

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

Isso tornará seus conjuntos de dados compatíveis com a API tfds.load / tfds.builder , simplesmente fornecendo data_dir/ :

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