A model that won’t learn or won’t generalize needs the right techniques: regularization against overfitting, schedules and normalization for stability, and the GPU for speed.
Why: a network with enough capacity will memorize the training set — dropout, weight decay, and early stopping keep it generalizing to new data. When: add them when validation loss rises while training loss keeps falling. Where: dropout randomly zeros activations during training only (hence train()/eval() matter), forcing robustness.
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 64), nn.ReLU(),
nn.Dropout(0.3), # zero 30% of activations while training
nn.Linear(64, 3),
)
# Weight decay (L2) via the optimizer:
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
# Early stopping: halt when val loss hasn't improved for N epochs.Why: batch normalization steadies and speeds up training, and a learning-rate schedule that decays over time helps the model settle into a good minimum. When: add BatchNorm between linear/conv and activation; use a scheduler to lower the LR as training progresses. Where: the learning rate is the single most important hyperparameter — too high diverges, too low crawls.
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 64), nn.BatchNorm1d(64), nn.ReLU(),
nn.Linear(64, 3),
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3)
# After each epoch: scheduler.step(val_loss)Why: moving the model and each batch to the GPU is what makes training practical, and mixed precision (fp16/bf16) roughly halves memory and speeds it up further on modern GPUs. When: send model and data to the device; enable autocast + GradScaler on capable hardware. Where: model and data must be on the same device or you get an error.
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
scaler = torch.cuda.amp.GradScaler() # mixed precision
for x, y in train_loader:
x, y = x.to(device), y.to(device) # data to the same device
optimizer.zero_grad()
with torch.autocast(device_type=device): # fp16/bf16 where safe
loss = criterion(model(x), y)
scaler.scale(loss).backward()
scaler.step(optimizer); scaler.update()