Quick Recap#
This chapter has guided you through the essentials of building, training, and evaluating neural networks in PyTorch. Here’s a quick recap of the key concepts to remember.
Data Management. PyTorch provides the classes
DatasetandDataLoaderto load and iterate over batches of data. They can be combined with the APItorchvision.transforms.v2to apply transformations to image data before it is fed into the network.Building Neural Networks. PyTorch makes it easy to define neural networks by subclassing
torch.nn.Moduleand implementing theforwardmethod using the layers available intorch.nn.Loss Function. PyTorch provides a wide range of loss functions that can be used to train a model. The choice of loss function depends on the task at hand.
For binary classification, the network should be trained with
nn.BCEWithLogitsLoss, which combines the sigmoid activation and the binary cross-entropy loss.For multi-class classification, the network should be trained with
nn.CrossEntropyLoss, which combines the softmax activation and the negative log-likelihood loss.For regression, the network should be trained with
nn.MSELoss.
Optimizer. PyTorch provides a wide range of optimizers that can be used to train a model, such as
SGD,Adam,RMSpropfrom thetorch.optimmodule. The choice of optimizer depends on the task at hand, butAdamorAdamWis generally a good choice.Training Neural Networks. The training loop in PyTorch typically involves iterating over batches of data using a DataLoader, computing predictions with the model, calculating the prediction error using a loss function, and updating the model parameters with an optimizer. It is a good practice to isolate the training logic in a separate function, so that it can be reused easily to experiment with different architectures and hyperparameters.
Overfitting. Neural networks tend to overfit the training data, i.e., they perform well on the training data but poorly on unseen data. It is a good practice to monitor the performance of the model on a separate validation set during training. This can be done quite easily by modifying the training loop to include an evaluation phase.
Evaluation. Implementing evaluation metrics from scratch can be tedious and error-prone. Third-party libraries like TorchEval, TorchMetrics, and scikit-learn provide tested metrics that integrate with a learning pipeline.
Pipeline checklist#
Before trusting an experiment, check the whole path.
Keep training, validation, and test roles separate; shuffle only training data.
Verify one batch:
input shape, dtype and range;
target shape, dtype and class values.
Confirm that the model returns logits with the expected shape.
Pair logits and targets with the correct loss function.
Try to overfit one small batch before running a long experiment.
Monitor training and validation curves.
Use the test set once, after all choices are final, and inspect errors in addition to a headline score.
Common symptoms#
Loss does not move: check the learning rate,
requires_grad, gradient clearing, target dtype, and whether the optimizer received the model parameters.Training improves but validation worsens: the model is overfitting; consider stopping earlier, collecting more data, reducing capacity, or adding regularization.
Device mismatch: move the model, inputs, and targets to the same device.
Suspiciously excellent validation: look for duplicates, shared subjects, preprocessing leakage, or accidental training on validation data.
Check your understanding#
Why would three linear layers without activations still behave like a single linear layer?
For ten-class classification, what should the output shape, target dtype, and loss function be?
Why is the checkpoint with the smallest training loss not necessarily the best checkpoint?
What should you investigate if a model cannot memorize one small batch?
An MLP flattens a 28×28 image. What spatial information does that discard, and how might a convolutional network use it?