Datasets documentation
Uso con PyTorch
Uso con PyTorch
Este documento es una introducción rápida al uso de datasets con PyTorch, con un enfoque particular en cómo obtener objetos torch.Tensor de nuestros datasets y cómo utilizar un DataLoader de PyTorch y un Dataset de Hugging Face con el mejor rendimiento.
Formato del conjunto de datos
Por defecto, los conjuntos de datos devuelven objetos de Python estándar: enteros, flotantes, cadenas, listas, etc.
Para obtener tensores de PyTorch en su lugar, puede establecer el formato del conjunto de datos como pytorch usando Dataset.with_format():
>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
[3, 4]])}Un objeto Dataset es un envoltorio de una tabla Arrow, que permite lecturas rápidas sin copia desde las matrices del conjunto de datos hacia tensores de PyTorch.
Para cargar los datos como tensores en una GPU, especifique el argumento device:
>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda:0')}Arreglos de N dimensiones
Si su conjunto de datos consiste en matrices de N dimensiones, verá que por defecto se consideran el mismo tensor si la forma es fija:
>>> 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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
[3, 4]])}>>> 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': [tensor([1, 2]), tensor([3])]}Sin embargo, esta lógica a menudo requiere comparaciones de forma y copias de datos lentas.
Para evitarlo, debes usar explícitamente el tipo de función Array y especificar la forma de tus 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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
[3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])}Otros tipos de características
Los datos de ClassLabel se convierten correctamente en 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("torch")
>>> ds[:3]
{'label': tensor([0, 0, 1])}Los objetos de cadena y binarios no se modifican, ya que PyTorch solo admite números.
Los tipos de características Image y Audio también son compatibles.
Para usar el tipo de función Image, deberá instalar el extra
visioncomopip 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("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]]], dtype=torch.uint8)}Para usar el tipo de función Audio, debes instalar el extra
audiocomopip 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("torch")
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05, 1.5259e-05, 1.6785e-04, ..., -1.5259e-05,
-1.5259e-05, 1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)Carga de datos
Al igual que los objetos torch.utils.data.Dataset, un Dataset se puede pasar directamente a un DataLoader de PyTorch:
>>> import numpy as np
>>> from datasets import Dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
... print(batch)
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}Optimizar la carga de datos
Hay varias formas de aumentar la velocidad de carga de tus datos, lo que puede ahorrarte tiempo, especialmente si trabajas con conjuntos de datos grandes. PyTorch ofrece carga de datos paralelizada, obteniendo lotes de índices en lugar de de forma individual, y streaming para iterar sobre el conjunto de datos sin descargarlo en el disco.
Usar múltiples Workers
Puede paralelizar la carga de datos con el argumento num_workers de un DataLoader de PyTorch y obtener un mayor rendimiento.
Por debajo, el DataLoader inicia num_workers procesos.
Cada proceso recarga el conjunto de datos pasado al DataLoader y se utiliza para consultar ejemplos.
Recargar el conjunto de datos dentro de un trabajador no llena tu RAM, ya que simplemente vuelve a mapear el conjunto de datos en memoria desde tu disco.
>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)Transmitir datos
Transmite un dataset cargándolo como un IterableDataset. Esto te permite iterar progresivamente sobre un dataset remoto sin descargarlo en disco y/o sobre archivos de datos locales. Aprende más sobre qué tipo de dataset es el más adecuado para tu caso de uso en la guía elección entre un dataset regular o un dataset iterable.
Un conjunto de datos iterable de datasets hereda de torch.utils.data.IterableDataset, por lo que puedes pasarlo a un torch.utils.data.DataLoader:
>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset") # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)Si el conjunto de datos se divide en varios fragmentos (es decir, si el conjunto de datos consta de varios archivos de datos), entonces se puede transmitir en paralelo usando num_workers:
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)En este caso, a cada worker se le asigna un subconjunto de la lista de shards desde los que transmitir.
Punto de control y reanudación
Si necesitas un DataLoader que puedas guardar y reanudar en medio del entrenamiento, puedes usar el StatefulDataLoader de torchdata:
>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)Esto es posible gracias a IterableDataset.state_dict() y IterableDataset.load_state_dict().
Distribuido
Para dividir tu conjunto de datos entre tus nodos de entrenamiento, puedes usar datasets.distributed.split_dataset_by_node():
import os
from datasets.distributed import split_dataset_by_node
ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))Esto funciona tanto para conjuntos de datos de estilo mapa como para conjuntos de datos iterables.
El conjunto de datos se divide para el nodo con rango rank en un grupo de nodos de tamaño world_size.
Para conjuntos de datos de estilo mapa:
A cada nodo se le asigna un fragmento de datos, por ejemplo, al rank 0 se le asigna el primer fragmento del conjunto de datos.
Para conjuntos de datos iterables:
La strategy de división puede ser "auto", "shards" o "examples":
"shards": los shards se asignan de manera uniforme entre los nodos"examples": cada nodo conserva 1 ejemplo deworld_size, omitiendo los demás ejemplos"auto"(predeterminado): usa"shards"si el conjunto de datos tiene un número de fragmentos que es un factor deworld_size(es decir, sidataset.num_shards % world_size == 0), que es la opción más optimizada. En caso contrario, usa"examples".
También se puede combinar con un torch.utils.data.DataLoader si se desea que cada nodo utilice varios workers para cargar los datos.
Update on GitHubSi barajas tu conjunto de datos iterable en una configuración distribuida, asegúrate de establecer una
seedfija en IterableDataset.shuffle() para que se utilice la misma lista barajada de fragmentos en cada nodo y se sepa qué fragmentos debe omitir el nodo.