Datasets & DataLoaders

PyTorch Dataset e DataLoader alimentam loops de treinamento com dados agrupados, embaralhados e opcionalmente aumentados. Pipelines de entrada eficientes mantêm a GPU ocupada e o treinamento rápido.

Receita

Cartão de referência rápida - pronto para copiar e colar.

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
 
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
for batch_x, batch_y in loader:
    batch_x = batch_x.to(device, non_blocking=True)

Quando usar isso:

• Carregando imagens, texto ou dados tabulares em loops de treinamento.
• Aplicando transformações por amostra (aumento, tokenização).
• Agrupando sequências de comprimento variável com um collate_fn personalizado.
• Pré-buscando dados para a GPU com pin_memory e non_blocking.

Exemplo de Trabalho

"""datasets_dataloaders.py - Dataset e DataLoader personalizados para dados CSV."""
from __future__ import annotations
 
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader, random_split
 
class TabularDataset(Dataset):
    def __init__(self, df: pd.DataFrame, feature_cols: list[str], target_col: str):
        self.X = torch.tensor(df[feature_cols].values, dtype=torch.float32)
        self.y = torch.tensor(df[target_col].values, dtype=torch.long)
 
    def __len__(self) -> int:
        return len(self.y)
 
    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.X[idx], self.y[idx]
 
df = pd.read_csv("train.csv")
dataset = TabularDataset(df, feature_cols=["f1", "f2", "f3"], target_col="label")
 
train_size = int(0.8 * len(dataset))
train_ds, val_ds = random_split(dataset, [train_size, len(dataset) - train_size])
 
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=256, shuffle=False, num_workers=2)
 
for epoch in range(3):
    for x_batch, y_batch in train_loader:
        # x_batch: (128, 3), y_batch: (128,)
        pass

O que isso demonstra:

• Dataset personalizado com __len__ e __getitem__.
• random_split para partição de treino/validação.
• DataLoader lida com agrupamento e embaralhamento.
• pin_memory=True acelera a transferência CPU-GPU quando CUDA está disponível.

Mergulho Profundo

Como Funciona

• Dataset.__getitem__(i) retorna uma amostra; DataLoader coleta batch_size amostras.
• shuffle=True randomiza a ordem a cada época.
• num_workers inicia subprocessos para carregamento paralelo de dados.
• collate_fn personaliza como as amostras são mescladas em um lote (preenchimento para sequências).
• IterableDataset transmite dados para fontes muito grandes ou online.

Opções do DataLoader

Notas de Python

from torchvision import datasets, transforms
 
# Dataset de imagem integrado com transforms
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_set = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)

Armadilhas

• Aumento de dados na validação - transformações aleatórias corrompem a avaliação. Correção: pipelines de transformações separados para treino/validação; validação recebe apenas transformações determinísticas.
• num_workers no Windows/notebooks - problemas de multiprocessamento no Jupyter. Correção: num_workers=0 em notebooks; usar scripts para carregamento paralelo.
• Não usar pin_memory com CUDA - cópias mais lentas de host para dispositivo. Correção: pin_memory=True e non_blocking=True em .to(device).
• Carregar o dataset inteiro na RAM - OOM (Out Of Memory) em datasets grandes. Correção: carregamento preguiçoso em __getitem__, arquivos com mapeamento de memória ou IterableDataset.
• Tipos de tensor inconsistentes - misturar float/int causa erros no modelo. Correção: converter explicitamente em __getitem__.
• Embaralhar dados de validação - desnecessário e torna as curvas de perda ruidosas. Correção: shuffle=False para loaders de validação/teste.

Alternativas

FAQs

def pad_collate(batch):
    xs, ys = zip(*batch)
    lengths = [len(x) for x in xs]
    padded = torch.nn.utils.rnn.pad_sequence(xs, batch_first=True)
    return padded, torch.tensor(ys), torch.tensor(lengths)
 
loader = DataLoader(ds, collate_fn=pad_collate)

DataLoader(ds, num_workers=4, persistent_workers=True)

generator = torch.Generator().manual_seed(42)
loader = DataLoader(ds, shuffle=True, generator=generator)

from torch.utils.data import WeightedRandomSampler
sampler = WeightedRandomSampler(weights, num_samples=len(weights))
loader = DataLoader(ds, sampler=sampler)

Relacionados

• Training Loops - consumindo lotes
• PyTorch Basics - tipos de tensor
• GPUs & Mixed Precision - pin_memory e non_blocking
• Transfer Learning & Fine-Tuning - transforms pré-treinados
• Distributed Training - DistributedSampler

Versões da Stack: Esta página foi escrita para Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, e uv 0.6+.