Assista palestras, sessões de produtos, workshops e muito mais na lista de reprodução do Google I / O See

Dicas de desempenho

Este documento fornece dicas de desempenho específicas do TFDS. Observe que o TFDS fornece conjuntos de dados comotf.data.Dataset s, portanto, o conselho do guia tf.data ainda se aplica.

Conjuntos de dados de referência

Use tfds.benchmark(ds) para comparar qualquer objetotf.data.Dataset .

Certifique-se de indicar batch_size= para normalizar os resultados (por exemplo, 100 iter / s -> 3200 ex / s). Isso funciona com qualquer iterável (por exemplo, tfds.benchmark(tfds.as_numpy(ds)) ).

ds = tfds.load('mnist', split='train').batch(32).prefetch()
# Display some benchmark statistics
tfds.benchmark(ds, batch_size=32)
# Second iteration is much faster, due to auto-caching
tfds.benchmark(ds, batch_size=32)

Conjuntos de dados pequenos (<GB)

Todos os conjuntos de dados TFDS armazenam os dados em disco no formato TFRecord . Para pequenos conjuntos de dados (por exemplo, Mnist, Cifar, ...), a leitura de .tfrecord pode adicionar sobrecarga significativa.

Como esses conjuntos de dados cabem na memória, é possível melhorar significativamente o desempenho armazenando em cache ou pré-carregando o conjunto de dados. Observe que o TFDS armazena automaticamente em cache pequenos conjuntos de dados (consulte a próxima seção para obter detalhes).

Cache do conjunto de dados

Aqui está um exemplo de um pipeline de dados que armazena em cache explicitamente o conjunto de dados após normalizar as imagens.

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label


ds, ds_info = tfds.load(
    'mnist',
    split='train',
    as_supervised=True,  # returns `(img, label)` instead of dict(image=, ...)
    with_info=True,
)
# Applying normalization before `ds.cache()` to re-use it.
# Note: Random transformations (e.g. images augmentations) should be applied
# after both `ds.cache()` (to avoid caching randomness) and `ds.batch()` (for
# vectorization [1]).
ds = ds.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds = ds.cache()
# For true randomness, we set the shuffle buffer to the full dataset size.
ds = ds.shuffle(ds_info.splits['train'].num_examples)
# Batch after shuffling to get unique batches at each epoch.
ds = ds.batch(128)
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)

Ao iterar sobre este conjunto de dados, a segunda iteração será muito mais rápida do que a primeira graças ao armazenamento em cache.

Cache automático

Por padrão, o TFDS armazena em cache automaticamente os conjuntos de dados que atendem às seguintes restrições:

  • O tamanho total do conjunto de dados (todas as divisões) é definido e <250 MiB
  • shuffle_files está desabilitado ou apenas um único fragmento é lido

É possível cancelar o cache automático passando try_autocaching=False para tfds.ReadConfig em tfds.load . Dê uma olhada na documentação do catálogo do conjunto de dados para ver se um conjunto de dados específico usará o cache automático.

Carregando os dados completos como um único Tensor

Se seu conjunto de dados couber na memória, você também pode carregar o conjunto de dados completo como um único Tensor ou array NumPy. É possível fazer isso definindo batch_size=-1 para batch_size=-1 todos os exemplos em um único tf.Tensor . Em seguida, use tfds.as_numpy para a conversão de tf.Tensor em np.array .

(img_train, label_train), (img_test, label_test) = tfds.as_numpy(tfds.load(
    'mnist',
    split=['train', 'test'],
    batch_size=-1,
    as_supervised=True,
))

Grandes conjuntos de dados

Grandes conjuntos de dados são fragmentados (divididos em vários arquivos) e normalmente não cabem na memória, portanto, não devem ser armazenados em cache.

Embaralhar e treinar

Durante o treinamento, é importante misturar bem os dados; dados mal embaralhados podem resultar em menor precisão de treinamento.

Além de usar ds.shuffle para embaralhar registros, você também deve definir shuffle_files=True para obter um bom comportamento de embaralhamento para conjuntos de dados maiores que são fragmentados em vários arquivos. Caso contrário, as épocas lerão os fragmentos na mesma ordem e, portanto, os dados não serão verdadeiramente randomizados.

ds = tfds.load('imagenet2012', split='train', shuffle_files=True)

Além disso, quando shuffle_files=True , TFDS desabilita options.experimental_deterministic , o que pode dar um ligeiro aumento de desempenho. Para obter embaralhamento determinístico, é possível cancelar esse recurso com tfds.ReadConfig : seja configurando read_config.shuffle_seed ou substituindo read_config.options.experimental_deterministic .

Fragmente automaticamente seus dados entre os trabalhadores

Quando o treinamento em vários trabalhadores, você pode usar o input_context argumento de tfds.ReadConfig , de modo que cada trabalhador irá ler um subconjunto dos dados.

input_context = tf.distribute.InputContext(
    input_pipeline_id=1,  # Worker id
    num_input_pipelines=4,  # Total number of workers
)
read_config = tfds.ReadConfig(
    input_context=input_context,
)
ds = tfds.load('dataset', split='train', read_config=read_config)

Isso é complementar à API subsplit. Primeiro, a API subplit é aplicada ( train[:50%] é convertido em uma lista de arquivos a serem lidos) e, em seguida, uma ds.shard() é aplicada a esses arquivos. Exemplo: ao usar train[:50%] com num_input_pipelines=2 , cada um dos 2 trabalhadores lerá 1/4 dos dados.

Quando shuffle_files=True , os arquivos são misturados em um trabalhador, mas não entre os trabalhadores. Cada trabalhador lerá o mesmo subconjunto de arquivos entre as épocas.

Descodificação de imagem mais rápida

Por padrão, o TFDS decodifica automaticamente as imagens. No entanto, há casos em que pode ser mais tfds.decode.SkipDecoding pular a decodificação da imagem com tfds.decode.SkipDecoding e aplicar manualmente o op tf.io.decode_image :

  • Ao filtrar exemplos (com ds.filter ), para decodificar imagens após os exemplos terem sido filtrados.
  • Ao cortar imagens, use o op tf.image.decode_and_crop_jpeg fundido.

O código para ambos os exemplos está disponível no guia de decodificação .