Python / PyTorch Fundamentals Interview Questions
What are Dataset and DataLoader in PyTorch and how do they work together?
PyTorch's data pipeline follows a clean two-class design: Dataset defines how to access a single sample (index → data), and DataLoader wraps a Dataset to handle batching, shuffling, and parallel loading.
import torch from torch.utils.data import Dataset, DataLoader import numpy as np class TabularDataset(Dataset): def __init__(self, X: np.ndarray, y: np.ndarray): # Convert once at construction â not inside __getitem__! self.X = torch.tensor(X, dtype=torch.float32) self.y = torch.tensor(y, dtype=torch.long) def __len__(self) -> int: """Required â tells DataLoader how many samples exist.""" return len(self.X) def __getitem__(self, idx: int): """Required â return a single (features, label) sample.""" return self.X[idx], self.y[idx] # Synthetic data X = np.random.randn(1000, 20).astype(np.float32) y = np.random.randint(0, 3, size=1000) dataset = TabularDataset(X, y) print(len(dataset)) # 1000 print(dataset[0]) # (tensor of 20 features, tensor scalar label) loader = DataLoader( dataset, batch_size=32, shuffle=True, # shuffle each epoch â essential for training num_workers=4, # parallel data loading subprocesses pin_memory=True, # faster CPUâGPU transfer drop_last=True, # drop incomplete final batch ) # Iterate over batches for X_batch, y_batch in loader: print(X_batch.shape, y_batch.shape) # (32, 20) (32,) break # torchvision pre-built datasets from torchvision import datasets, transforms transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]) mnist = datasets.MNIST(root="./data", train=True, download=True, transform=transform) mnist_loader = DataLoader(mnist, batch_size=64, shuffle=True)
| Component | Responsibility | Required methods |
|---|---|---|
| Dataset | Defines how to access ONE sample by index | __len__, __getitem__ |
| DataLoader | Batches samples, shuffles, parallelises loading | Wraps any Dataset object |
More Related questions...