Datasets documentation
Usar com JAX
Usar com JAX
Este documento é uma introdução rápida ao uso de datasets com JAX, com foco especial em como obter objetos jax.Array dos nossos datasets e como usá-los para treinar modelos JAX.
jaxejaxlibsão necessários para reproduzir o código acima, portanto, certifique-se de instalá-los comopip install datasets[jax].
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., e objetos de string e binários permanecem inalterados, pois o JAX suporta apenas números.
Para obter arrays do JAX (estilo numpy) em vez disso, você pode definir o formato do dataset como jax:
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Um objeto Dataset é um encapsulamento de uma tabela Arrow, que permite leituras rápidas de arrays no dataset para arrays JAX.
Observe que o mesmo procedimento exato se aplica a objetos DatasetDict, de modo que, ao definir o formato de um DatasetDict como jax, todos os Datasets presentes nele
serão formatados como jax:
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("jax")
>>> dds["train"][:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Outra coisa que você precisa levar em consideração é que o formatação não é aplicada até que você realmente acesse os dados. Portanto, se você quiser obter um array JAX de um dataset, precisará acessar os dados primeiro; caso contrário, o formato permanecerá o mesmo.
Por fim, para carregar os dados no dispositivo de sua escolha, você pode especificar o argumento device, mas observe que jaxlib.xla_extension.Device não é suportado, pois não é serializável nem com pickle nem com dill, então você precisará usar seu identificador em string no lugar:
>>> import jax
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> device = str(jax.devices()[0]) # Not casting to `str` before passing it to `with_format` will raise a `ValueError`
>>> ds = ds.with_format("jax", device=device)
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[0]["data"].device()
TFRT_CPU_0
>>> assert ds[0]["data"].device() == jax.devices()[0]
TrueObserve que, se o argumento device não for fornecido para with_format, ele usará o dispositivo padrão, que é jax.devices()[0].
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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}>>> 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("jax")
>>> ds[0]
{'data': [Array([1, 2], dtype=int32), Array([3], dtype=int32)]}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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}
>>> ds[:2]
{'data': Array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]], dtype=int32)}Outros tipos de recursos
Os dados de ClassLabel são convertidos corretamente em arrays:
>>> 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("jax")
>>> ds[:3]
{'label': DeviceArray([0, 0, 1], dtype=int32)}Objetos de string e binários permanecem inalterados, pois o JAX 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
visioncomopip install datasets[vision].
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': DeviceArray([[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': DeviceArray([[[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]]], dtype=uint8)}Para usar o tipo de recurso Audio, você precisará instalar o extra
audiocomopip install datasets[audio].
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["audio"]["array"]
DeviceArray([-0.059021 , -0.03894043, -0.00735474, ..., 0.0133667 ,
0.01809692, 0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
DeviceArray(44100, dtype=int32, weak_type=True)Carregamento de dados
O JAX não possui recursos nativos de carregamento de dados, portanto, você precisará usar uma biblioteca como o PyTorch para carregar seus dados usando um DataLoader ou o TensorFlow usando um tf.data.Dataset. Citando a documentação do JAX sobre esse tópico:
“O JAX é focado em transformações de programas e NumPy com suporte de aceleradores, por isso não incluímos o carregamento ou a manipulação de dados na biblioteca JAX. Já existem muitos carregadores de dados excelentes por aí, então vamos apenas usá-los em vez de reinventar a roda. Vamos pegar o carregador de dados do PyTorch e criar um pequeno shim para fazê-lo funcionar com arrays do NumPy.”.
É por isso que o formato JAX em datasets é tão útil, pois permite usar qualquer modelo do HuggingFace Hub com JAX, sem precisar se preocupar com a parte de carregamento de dados.
Usando with_format('jax')
A maneira mais fácil de obter arrays do JAX de um dataset é usar o método with_format('jax'). Vamos supor que queremos treinar uma rede neural no dataset MNIST disponível no HuggingFace Hub em https://huggingface.co/datasets/ylecun/mnist.
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("jax")
>>> ds["train"][0]
{'image': DeviceArray([[ 0, 0, 0, ...],
[ 0, 0, 0, ...],
...,
[ 0, 0, 0, ...],
[ 0, 0, 0, ...]], dtype=uint8),
'label': DeviceArray(5, dtype=int32)}Uma vez que o formato é definido, podemos alimentar o modelo JAX com o conjunto de dados em lotes usando o método Dataset.iter():
>>> for epoch in range(epochs):
... for batch in ds["train"].iter(batch_size=32):
... x, y = batch["image"], batch["label"]
... ...