Use com PyTorch

Este documento é uma introdução rápida ao uso de datasets com PyTorch, com foco especial em como obter objetos torch.Tensor dos nossos datasets e como usar um DataLoader do PyTorch e um Dataset do Hugging Face com o melhor desempenho.

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 PyTorch, você pode definir o formato do conjunto de dados 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]])}

Um objeto Dataset é um encapsulamento de uma tabela Arrow, que permite leituras rápidas sem cópia de dados dos arrays no dataset para tensores do PyTorch.

Para carregar os dados como tensores em uma GPU, especifique o 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')}

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("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])]}

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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}
>>> ds[:2]
{'data': tensor([[[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("torch")  
>>> ds[:3]
{'label': tensor([0, 0, 1])}

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("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 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("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)

Carregamento de dados

Assim como objetos torch.utils.data.Dataset, um Dataset pode ser passado diretamente para um DataLoader do 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])}

Otimizar o carregamento de dados

Existem várias maneiras de aumentar a velocidade de carregamento dos seus dados, o que pode economizar tempo, especialmente se você estiver trabalhando com grandes conjuntos de dados. O PyTorch oferece carregamento de dados paralelizado, recuperando lotes de índices em vez de individualmente, e streaming para iterar sobre o conjunto de dados sem baixá-lo no disco.

Use múltiplos Workers

Você pode paralelizar o carregamento de dados com o argumento num_workers de um DataLoader do PyTorch e obter um throughput maior.

Por baixo dos panos, o DataLoader inicia num_workers processos. Cada processo recarrega o conjunto de dados passado para o DataLoader e é usado para consultar exemplos. Recarregar o conjunto de dados dentro de um worker não enche sua RAM, pois ele simplesmente mapeia o conjunto de dados novamente a partir do seu 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 dados

Faça o streaming de um dataset carregando-o como um IterableDataset. Isso permite iterar progressivamente sobre um dataset remoto sem baixá-lo no disco e/ou sobre arquivos de dados locais. Saiba mais sobre qual tipo de dataset é o ideal para o seu caso de uso no guia escolhendo entre um dataset regular ou um dataset iterável.

Um dataset iterável do datasets herda de torch.utils.data.IterableDataset, então você pode passá-lo para um 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)

Se o conjunto de dados for dividido em vários shards (ou seja, se o conjunto de dados consistir em múltiplos arquivos de dados), então você pode fazer streaming em 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)

Nesse caso, cada worker recebe um subconjunto da lista de shards para transmitir.

Checkpoint e retomada

Se você precisa de um DataLoader que possa ser checkpointado e retomado no meio do treinamento, você pode usar o StatefulDataLoader do 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)

Isso é possível graças a IterableDataset.state_dict() e IterableDataset.load_state_dict().

Distribuído

Para dividir seu conjunto de dados entre seus nós de treinamento, você pode 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"]))

Isso funciona para conjuntos de dados no estilo mapa e conjuntos de dados iteráveis. O conjunto de dados é dividido para o nó no rank rank em um pool de nós de tamanho world_size.

Para conjuntos de dados no estilo mapa:

Cada nó recebe um bloco de dados, por exemplo, o rank 0 recebe o primeiro bloco do conjunto de dados.

Para conjuntos de dados iteráveis:

A strategy de divisão pode ser "auto", "shards" ou "examples":

• "shards": os shards são distribuídos igualmente entre os nós
• "examples": cada nó mantém 1 exemplo de world_size, pulando os outros exemplos
• "auto" (padrão): usa "shards" se o conjunto de dados tiver um número de shards que é um fator de world_size (ou seja, se dataset.num_shards % world_size == 0), o que é o mais otimizado. Caso contrário, usa "examples".

Isso também pode ser combinado com um torch.utils.data.DataLoader se você quiser que cada nó use múltiplos workers para carregar os dados.

Se você embaralhar seu conjunto de dados iterável em uma configuração distribuída, certifique-se de definir uma seed fixa em IterableDataset.shuffle() para que a mesma lista embaralhada de shards seja usada em todos os nós, a fim de saber quais shards o nó deve ignorar.

Paralelismo de dados fragmentado híbrido (HSDP)

A torch.distributed.device_mesh.DeviceMesh descreve sua topologia de treinamento como uma grade de ranks, com uma dimensão para cada estratégia de paralelismo. Apenas as dimensões de paralelismo de dados importam para o carregamento de dados, e o sharding híbrido as torna fáceis de serem ignoradas:

• "fsdp" (geralmente intra-nó): o grupo FSDP em que cada rank mantém apenas um shard dos estados do modelo (parâmetros, gradientes e estados do otimizador) e faz temporariamente um all-gather dos parâmetros completos de um módulo logo antes de calcular sobre ele. O treinamento permanece paralelo de dados: cada rank executa um passo completo de forward e backward no seu próprio batch, então essa dimensão precisa de dados diferentes por rank, assim como no DDP.
• "dp_replicate" (geralmente entre nós): as cópias duplicadas desse modelo particionado, mantidas sincronizadas por um all-reduce adicional entre as réplicas. As réplicas também processam dados diferentes.
• "tp" (intra-nó): ranks de paralelismo de tensor que dividem os pesos de cada camada para calcular um lote juntos. Eles não são paralelos de dados e devem ler os mesmos exemplos.

O grau de paralelismo de dados é o produto das dimensões de sharding (FSDP) e de replicação, que é menor que WORLD_SIZE assim que a malha tiver uma dimensão como "tp". Passe para datasets.distributed.split_dataset_by_node() esse produto e o rank achatado sobre essas duas dimensões:

from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node

# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]

dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()

ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)

Aqui, dp_size é 2 * 2 = 4: cada rank de paralelismo de dados lê um quarto do conjunto de dados, e os dois ranks de cada par "tp" leem os mesmos exemplos.