Training Loop, Datasets & DataLoaders

The standard PyTorch workflow#

Almost every PyTorch project follows the same pattern:

1. Wrap data in a Dataset.
2. Feed it in batches with a DataLoader.
3. Define a model, a loss function and an optimizer.
4. Run the training loop for several epochs.
5. Check performance on a validation set after each epoch.

In this lesson you'll build a complete example: classify 2D points as inside or outside a circle.

Step 1 — a custom Dataset#

A Dataset needs two methods: __len__ (how many examples) and __getitem__ (return one example).

import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset, random_split

torch.manual_seed(42)


class CircleDataset(Dataset):
    """2D points labelled 1 if they lie inside a circle of radius 1, else 0."""

    def __init__(self, n: int = 2000):
        self.X = torch.rand(n, 2) * 4 - 2                       # points in [-2, 2]
        self.y = ((self.X ** 2).sum(dim=1) < 1).float().unsqueeze(1)

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

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

For real projects __getitem__ might load an image from disk or tokenise a sentence. Data is loaded lazily, one example at a time, so datasets larger than memory still work.

Step 2 — split and batch with DataLoader#

dataset = CircleDataset()
train_set, val_set = random_split(dataset, [1600, 400])

train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256)

• batch_size — how many examples per gradient update.
• shuffle=True — new order every epoch for training (never needed for validation).
• num_workers — load batches in background processes for faster training on big datasets.

Step 3 — model, loss and optimizer#

model = nn.Sequential(
    nn.Linear(2, 32), nn.ReLU(),
    nn.Linear(32, 32), nn.ReLU(),
    nn.Linear(32, 1),
)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

Step 4 — the validation function#

def evaluate():
    model.eval()                                  # dropout/batchnorm in eval mode
    correct, total = 0, 0
    with torch.inference_mode():                  # no gradients needed
        for X, y in val_loader:
            X, y = X.to(device), y.to(device)
            preds = (torch.sigmoid(model(X)) > 0.5).float()
            correct += (preds == y).sum().item()
            total += len(y)
    return correct / total

Step 5 — the training loop#

for epoch in range(1, 21):
    model.train()
    running = 0.0
    for X, y in train_loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()            # 1. clear old gradients
        loss = loss_fn(model(X), y)      # 2. forward pass + loss
        loss.backward()                  # 3. backpropagation
        optimizer.step()                 # 4. update weights
        running += loss.item() * len(X)

    if epoch % 5 == 0:
        print(f"epoch {epoch:2d}  train loss {running / len(train_set):.4f}  val acc {evaluate():.3f}")

Output on a CPU with PyTorch 2.15 (numbers can differ slightly on other versions or hardware):

epoch  5  train loss 0.0443  val acc 0.993
epoch 10  train loss 0.0289  val acc 0.973
epoch 15  train loss 0.0314  val acc 0.980
epoch 20  train loss 0.0253  val acc 0.990

The model reaches about 99% validation accuracy. Notice that validation accuracy moves up and down between checks — that's normal noise from mini-batch training, which is why you track the trend and keep the best checkpoint rather than trusting the last epoch.

Memorise the four lines inside the batch loop — zero_grad → forward + loss → backward → step. They appear in virtually every PyTorch training script.

What to watch during training#

Useful additions for real projects#

• Learning-rate schedulers: torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) — call scheduler.step() once per epoch.
• Gradient clipping: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) before optimizer.step() to prevent exploding gradients.
• Mixed precision with torch.autocast to train faster on modern GPUs.
• Early stopping: keep the weights from the best validation epoch.

Common mistakes#

• Forgetting optimizer.zero_grad() — gradients accumulate across batches.
• Data and model on different devices — move every batch with .to(device).
• Shuffling the validation set or computing gradients during validation.
• Using loss (a tensor) for logging — use loss.item() so you don't keep the graph in memory.

Key takeaways#

• Dataset returns single examples; DataLoader batches and shuffles them.
• The core loop: zero_grad(), forward + loss, backward(), step().
• Validate every epoch in eval() mode under inference_mode().
• Watch train vs validation metrics to catch overfitting early.