# 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()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.with_format):

```py
>>> 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]])}
```

> [!TIP]
> Um objeto [Dataset](/docs/datasets/main/pt/package_reference/main_classes#datasets.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`:
```py
>>> 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:

```py
>>> 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]])}
```

```py
>>> 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 tensors:

```py
>>> 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](/docs/datasets/main/pt/package_reference/main_classes#datasets.ClassLabel) são convertidos corretamente para tensores:

```py
>>> 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](/docs/datasets/main/pt/package_reference/main_classes#datasets.Image) e [Audio](/docs/datasets/main/pt/package_reference/main_classes#datasets.Audio) também são suportados.

> [!TIP]
> Para usar o tipo de recurso [Image](/docs/datasets/main/pt/package_reference/main_classes#datasets.Image), você precisará instalar o extra `vision` como `pip install datasets[vision]`.

```py
>>> 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)}
```

> [!TIP]
> Para usar o tipo de recurso [Audio](/docs/datasets/main/pt/package_reference/main_classes#datasets.Audio), você precisa instalar o extra `audio` como `pip install datasets[audio]`.

```py
>>> 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](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset) pode ser passado diretamente para um `DataLoader` do PyTorch:

```py
>>> 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, a recuperação de lotes de índices em vez de individualmente e o 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.

```py
>>> 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](/docs/datasets/main/pt/package_reference/main_classes#datasets.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](./about_mapstyle_vs_iterable).

Um dataset iterável do `datasets` herda de `torch.utils.data.IterableDataset`, então você pode passá-lo para um `torch.utils.data.DataLoader`:

```py
>>> 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`:

```py
>>> 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](https://github.com/pytorch/data):

```py
>>> 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()](/docs/datasets/main/pt/package_reference/main_classes#datasets.IterableDataset.state_dict) e [IterableDataset.load_state_dict()](/docs/datasets/main/pt/package_reference/main_classes#datasets.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()](/docs/datasets/main/pt/package_reference/main_classes#datasets.distributed.split_dataset_by_node):

```python
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ó com 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.

> [!WARNING]
> 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()](/docs/datasets/main/pt/package_reference/main_classes#datasets.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`](https://pytorch.org/docs/stable/distributed.html#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](https://pytorch.org/docs/stable/fsdp.html) 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()](/docs/datasets/main/pt/package_reference/main_classes#datasets.distributed.split_dataset_by_node) esse produto e o rank achatado sobre essas duas dimensões:

```python
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.