Usando Datasets com TensorFlow

Este documento é uma introdução rápida ao uso de datasets com TensorFlow, com foco especial em como obter objetos tf.Tensor dos nossos datasets e como transmitir dados de objetos Dataset do Hugging Face para métodos do Keras, como model.fit().

Formato do conjunto de dados

Por padrão, os conjuntos de dados retornam objetos Python comuns: inteiros, números de ponto flutuante, strings, listas, etc.

Para obter tensores do TensorFlow, você pode definir o formato do conjunto de dados como tf:

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

Um objeto Dataset é um encapsulamento de uma tabela Arrow, que permite leituras rápidas de arrays no dataset para tensores do TensorFlow.

Isso pode ser útil para converter seu conjunto de dados em um dicionário de objetos Tensor, ou para escrever um gerador para carregar amostras do TF a partir dele. Se você deseja converter todo o conjunto de dados para Tensor, basta consultar o conjunto de dados completo:

>>> ds[:]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

Arrays N-dimensionais

Se o seu conjunto de dados consiste em arrays N-dimensionais, você verá que, por padrão, eles são considerados o mesmo tensor se a forma for fixa:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}

Caso contrário, um dataset formatado em TensorFlow produz um RaggedTensor em vez de um tensor único:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': <tf.RaggedTensor [[1, 2], [3]]>}

No entanto, essa lógica frequentemente exige comparações de forma e cópias de dados lentas. Para evitar isso, você deve usar explicitamente o tipo de recurso Array e especificar a forma dos seus tensores:

>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2, 2), dtype=int64, numpy=
 array([[[1, 2],
         [3, 4]],
 
        [[5, 6],
         [7, 8]]])>}

Outros tipos de recursos

Os dados de ClassLabel são convertidos corretamente para tensores:

>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[:3]
{'label': <tf.Tensor: shape=(3,), dtype=int64, numpy=array([0, 0, 1])>}

Strings e objetos binários também são suportados:

>>> from datasets import Dataset, Features 
>>> text = ["foo", "bar"]
>>> data = [0, 1] 
>>> ds = Dataset.from_dict({"text": text, "data": data})  
>>> ds = ds.with_format("tf") 
>>> ds[:2]
{'text': <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'foo', b'bar'], dtype=object)>,
 'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>}

Você também pode formatar explicitamente certas colunas e deixar as outras colunas sem formatação:

>>> ds = ds.with_format("tf", columns=["data"], output_all_columns=True)
>>> ds[:2]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>,
 'text': ['foo', 'bar']}

Objetos de string e binários permanecem inalterados, pois o PyTorch suporta apenas números.

Os tipos de recurso Image e Audio também são suportados.

Para usar o tipo de recurso Image, você precisará instalar o extra vision como pip install datasets[vision].

>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]
{'image': <tf.Tensor: shape=(512, 512, 4), dtype=uint8, numpy=
 array([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=uint8)>}
>>> ds[:2]
{'image': <tf.Tensor: shape=(2, 512, 512, 4), dtype=uint8, numpy=
 array([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=uint8)>}

Para usar o tipo de recurso Audio, você precisará instalar o extra audio como pip install datasets[audio].

>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]["audio"]["array"]
<tf.Tensor: shape=(202311,), dtype=float32, numpy=
array([ 6.1035156e-05,  1.5258789e-05,  1.6784668e-04, ...,
       -1.5258789e-05, -1.5258789e-05,  1.5258789e-05], dtype=float32)>
>>> ds[0]["audio"]["sampling_rate"]
<tf.Tensor: shape=(), dtype=int32, numpy=44100>

Carregamento de dados

Embora você possa carregar amostras e lotes individuais apenas indexando seu dataset, isso não funcionará se você quiser usar métodos do Keras como fit() e predict(). Você poderia escrever uma função geradora que embaralha e carrega lotes do seu dataset e usar fit() nela, mas isso parece ser um trabalho desnecessário. Em vez disso, se você quiser transmitir dados do seu dataset em tempo real, recomendamos converter seu dataset em um tf.data.Dataset usando o método to_tf_dataset().

A classe tf.data.Dataset abrange uma ampla variedade de casos de uso - ela é frequentemente criada a partir de Tensors em memória, ou usando uma função de carregamento para ler arquivos em disco ou armazenamento externo. O dataset pode ser transformado arbitrariamente com o método map(), ou métodos como batch() e shuffle() podem ser usados para criar um dataset pronto para treinamento. Esses métodos não modificam os dados armazenados de nenhuma forma - em vez disso, os métodos constroem um grafo de pipeline de dados que será executado quando o dataset for iterado, geralmente durante o treinamento ou inferência do modelo. Isso é diferente do método map() dos objetos Dataset do Hugging Face, que executa a função de mapeamento imediatamente e salva as novas ou alteradas colunas.

Como todo o pipeline de pré-processamento de dados pode ser compilado em um tf.data.Dataset, essa abordagem permite o carregamento e o treinamento de dados de forma massivamente paralela e assíncrona. No entanto, o requisito de compilação de grafo pode ser uma limitação, particularmente para os tokenizers da Hugging Face, que geralmente não são (ainda!) compiláveis como parte de um grafo TF. Como resultado, geralmente recomendamos pré-processar o dataset como um dataset da Hugging Face, onde funções Python arbitrárias podem ser usadas, e depois converter para tf.data.Dataset em seguida usando to_tf_dataset() para obter um dataset em lotes pronto para o treinamento. Para ver exemplos dessa abordagem, consulte os exemplos ou notebooks para transformers.

Usando to_tf_dataset()

O uso de to_tf_dataset() é direto. Uma vez que seu conjunto de dados estiver pré-processado e pronto, basta chamá-lo da seguinte forma:

>>> from datasets import Dataset
>>> data = {"inputs": [[1, 2],[3, 4]], "labels": [0, 1]}
>>> ds = Dataset.from_dict(data)
>>> tf_ds = ds.to_tf_dataset(
            columns=["inputs"],
            label_cols=["labels"],
            batch_size=2,
            shuffle=True
            )

O objeto tf_ds retornado aqui está agora totalmente pronto para o treinamento e pode ser passado diretamente para model.fit(). Observe que você define o tamanho do lote ao criar o conjunto de dados, portanto, não precisa especificá-lo ao chamar fit():

>>> model.fit(tf_ds, epochs=2)

Para uma descrição completa dos argumentos, consulte a documentação de to_tf_dataset(). Em muitos casos, você também precisará adicionar um collate_fn à sua chamada. Esta é uma função que recebe múltiplos elementos do conjunto de dados e os combina em um único lote. Quando todos os elementos têm o mesmo comprimento, o colador padrão embutido será suficiente, mas para tarefas mais complexas, um colador personalizado pode ser necessário. Em particular, muitas tarefas possuem amostras com comprimentos de sequência variáveis, o que exigirá um data collator que possa preencher os lotes corretamente. Você pode ver exemplos disso nos exemplos e notebooks de NLP do transformers, onde comprimentos de sequência variáveis são muito comuns.

Se você perceber que o carregamento com to_tf_dataset está lento, também pode usar o argumento num_workers. Isso inicia vários subprocessos para carregar os dados em paralelo. Esse recurso é recente e ainda é um pouco experimental - por favor, abra uma issue se encontrar algum bug ao usá-lo!

Quando usar to_tf_dataset

O leitor atento pode ter notado até aqui que oferecemos duas abordagens para alcançar o mesmo objetivo - se você quiser passar seu conjunto de dados para um modelo do TensorFlow, você pode converter o conjunto de dados para um Tensor ou dict de Tensors usando .with_format('tf'), ou pode converter o conjunto de dados para um tf.data.Dataset com to_tf_dataset(). Qualquer uma dessas pode ser passada para model.fit(), então qual você deve escolher?

O ponto-chave a ser reconhecido é que, ao converter o conjunto de dados inteiro para Tensors, ele se torna estático e é totalmente carregado na RAM. Isso é simples e conveniente, mas, se qualquer uma das seguintes condições se aplicar, você provavelmente deve usar to_tf_dataset() em vez disso:

• Seu conjunto de dados é grande demais para caber na RAM. to_tf_dataset() transmite apenas um lote por vez, portanto, mesmo conjuntos de dados muito grandes podem ser processados com este método.
• Você deseja aplicar transformações aleatórias usando dataset.with_transform() ou o collate_fn. Isso é comum em várias modalidades, como aumentos de imagem ao treinar modelos de visão ou mascaramento aleatório ao treinar modelos de linguagem mascarados. Usar to_tf_dataset() aplicará essas transformações no momento em que um lote é carregado, o que significa que as mesmas amostras receberão diferentes aumentos cada vez que forem carregadas. Isso geralmente é o que você deseja.
• Seus dados possuem uma dimensão variável, como os textos de entrada em PLN que consistem em números variados de tokens. Ao criar um lote com amostras de dimensão variável, a solução padrão é preencher as amostras mais curtas até o comprimento da mais longa. Ao transmitir amostras de um conjunto de dados com to_tf_dataset, você pode aplicar esse preenchimento a cada lote por meio do seu collate_fn. No entanto, se você deseja converter um conjunto de dados assim em Tensors densos, então terá que preencher as amostras até o comprimento da amostra mais longa em todo o conjunto de dados! Isso pode resultar em grandes quantidades de preenchimento, o que desperdiça memória e reduz a velocidade do seu modelo.

Observações e limitações

No momento, to_tf_dataset() sempre retorna um dataset em lotes - em breve, adicionaremos suporte para datasets sem lotes!