Dataset and DataLoader

Code

import torch
from torch.utils.data import Dataset, DataLoader, TensorDataset

class MyDataset(Dataset):
    def __init__(self, x, y):
        self.x = torch.tensor(x, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

ds = MyDataset([[0, 0], [0, 1], [1, 0], [1, 1]], [0, 1, 1, 0])
loader = DataLoader(ds, batch_size=2, shuffle=True, num_workers=0, drop_last=False)

# Quick tensor dataset shortcut
xs = torch.randn(100, 3)
ys = torch.randint(0, 2, (100,))
quick = TensorDataset(xs, ys)
quick_loader = DataLoader(quick, batch_size=16, shuffle=True)

for xb, yb in loader:
    print(xb.shape, yb.shape)

Explanation

A Dataset subclasses torch.utils.data.Dataset and implements __len__ and __getitem__ to return one sample at a time. DataLoader batches, shuffles, and optionally parallelizes loading with num_workers. TensorDataset is a convenient shortcut when data already lives in tensors.

More PyTorch Snippets