Real training data won’t fit in memory and must be fed in batches. Dataset defines how to fetch one sample; DataLoader batches, shuffles, and parallelizes them.
Why: you cannot pass a whole dataset to a model at once — DataLoader serves it in shuffled batches, which is both memory-feasible and better for training. When: wrap any data source in a Dataset, then a DataLoader. Where: a Dataset says how to get one item; the DataLoader handles batching, shuffling, and parallel loading.
import torch
from torch.utils.data import TensorDataset, DataLoader
X = torch.randn(1000, 10) # 1000 samples, 10 features
y = torch.randint(0, 3, (1000,)) # 3 classes
dataset = TensorDataset(X, y)
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
for xb, yb in train_loader:
print(xb.shape, yb.shape) # torch.Size([32, 10]) torch.Size([32])
breakWhy: real data (images on disk, rows in a file) needs custom loading logic — you subclass Dataset and implement __len__ and __getitem__, and the DataLoader does the rest. When: write one whenever your data is not already tensors in memory. Where: __getitem__ loads and transforms a single sample, keeping memory use flat no matter the dataset size.
from torch.utils.data import Dataset
import torch
class CSVDataset(Dataset):
def __init__(self, features, labels):
self.X = torch.tensor(features, dtype=torch.float32)
self.y = torch.tensor(labels, dtype=torch.long)
def __len__(self):
return len(self.y)
def __getitem__(self, idx):
return self.X[idx], self.y[idx] # one sample at a time
loader = DataLoader(CSVDataset(feat, lab), batch_size=64, shuffle=True)Why: batch size trades memory for gradient stability, shuffling prevents the model from learning data order, and worker processes keep the GPU fed instead of waiting on data. When: shuffle the training set (never the test set), tune batch size to your memory, add workers if data loading is the bottleneck. Where: a starved GPU sitting idle usually means too few workers.
train_loader = DataLoader(
dataset,
batch_size=64, # bigger = more stable gradients, more memory
shuffle=True, # shuffle TRAIN only, so order isn't learned
num_workers=4, # parallel loading so the GPU never waits
pin_memory=True, # faster host->GPU transfer
)
# Validation/test loaders: shuffle=False.