Deep Learning Daily

How to Build Your First Image Classifier in PyTorch: A Step-by-Step Tutorial

How to Build Your First Image Classifier in PyTorch: A Step-by-Step Tutorial

  • September 20, 2026

Fashion-MNIST is a small dataset: 70,000 grayscale images of clothing, each 28 by 28 pixels, sorted into ten categories such as T-shirt, trouser, and sneaker. It is boring in the best possible way. Big enough to teach a real lesson, small enough that a laptop finishes training before your coffee goes cold. By the end of this walkthrough you will have a working classifier and, more usefully, a mental model of the four pieces every training script has: data, network, loss, and loop.

Set Up the Environment

You need Python plus two libraries: torch and torchvision. A single command installs both:

pip install torch torchvision

Then check what you are working with. Print torch.__version__, and call torch.cuda.is_available(). If that returns False, keep going anyway. This model trains on a CPU in a few minutes, and you will learn more by finishing a run than by configuring hardware.

Load and Inspect the Data

One line gets you the training split: torchvision.datasets.FashionMNIST(root="./data", train=True, download=True, transform=ToTensor()). Setting train=False gives you the test split instead. The transform matters: ToTensor converts each image into a tensor and rescales pixel values from 0–255 down to 0–1, which keeps your early gradients from being enormous.

Raw datasets are awkward to train on directly. A DataLoader handles batching and shuffling for you: DataLoader(train_set, batch_size=64, shuffle=True). Batching keeps memory use predictable. Shuffling stops the model from learning the order of your files rather than the content of your images. For the test set, set shuffle=False.

Sanity-check one batch first

Pull a single batch and print the shapes before you train anything. The images should come out as [64, 1, 28, 28] — sixty-four images, one colour channel, 28 by 28 — and the labels as [64]. If the shapes are wrong, everything downstream is wrong too, and you want to discover that in two seconds rather than after a failed run. Plot a few samples with matplotlib (call squeeze() on the channel dimension, then imshow) and check that the labels match what you see.

  • The batch dimension matches your chosen batch_size.
  • Pixel values sit between 0 and 1.
  • Labels are integers from 0 to 9.
  • Two consecutive shuffled batches are not identical.

Build the Network

A classifier is an nn.Module subclass with two methods. In __init__ you declare the layers. In forward you describe how data flows through them.

Ours is deliberately plain. Flatten each image into a vector of 784 numbers (28 × 28), send it through a Linear layer with 128 outputs, apply ReLU to keep the useful signals and zero the rest, then finish with a second Linear layer producing 10 scores — one per class. Those raw scores are called logits.

Two details that trip up beginners

Do not add a softmax at the end. nn.CrossEntropyLoss applies one internally and expects raw logits, so stacking your own on top makes the loss look reasonable while the gradients quietly misbehave.

Declare layers in __init__, never inside forward. Layers created during the forward pass are rebuilt on every iteration, so the optimizer ends up updating weights that no longer exist.

Choose a Loss Function and an Optimizer

Two lines cover it. criterion = nn.CrossEntropyLoss() measures how wrong the predictions are. optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) decides how the weights change. Adam is a forgiving default because it adapts the step size per parameter, which means you do not have to tune the learning rate on your first run. Plain SGD with momentum 0.9 works too, but it asks more of you.

The critical part is model.parameters(). Hand the optimizer anything else and training will run without a single error while the weights never move.

Write the Training Loop

An epoch is one full pass over the training set. Five to ten epochs is plenty for this model. Inside each epoch, loop over the DataLoader and repeat these steps for every batch:

  1. Move the batch to your device, if you are using one.
  2. Clear old gradients with optimizer.zero_grad().
  3. Run the forward pass: outputs = model(images).
  4. Compute the loss: loss = criterion(outputs, labels).
  5. Run loss.backward() to calculate gradients.
  6. Update the weights with optimizer.step().

Collect the loss as you go and print the average at the end of each epoch. Averages are easier to compare across epochs than sums, and a number that falls steadily is the clearest sign that the wiring is correct.

If the loss has not moved after the first epoch, check the order of zero_grad, backward, and step before you touch the architecture.

Measure Accuracy on Unseen Data

Accuracy on your training data tells you how well the model memorized. Accuracy on the test set tells you whether it learned anything. Wrap the evaluation in model.eval() and torch.no_grad(). The first switches off dropout and batch normalization behaviour — this network has neither, but the habit pays off the moment you add them. The second stops PyTorch from building a computation graph you will never use, which saves both memory and time.

For each test batch, take argmax along the class dimension to get the predicted label, compare it against the true labels, and count the matches. Divide the total number correct by the number of test images.

Expect a result in the high eighties. If you land near 10%, the model is guessing across ten classes — check that you are feeding it images rather than labels. If training accuracy is 99% and test accuracy is 70%, you are overfitting: shrink the hidden layer, add dropout, or accept that a fully connected network has limits on image data.

Where to Go Next

Keep the loop and swap one piece at a time. That is how the concepts actually stick.

  • Save your weights with torch.save(model.state_dict(), "fashion_mlp.pt") so you can reload them later.
  • Switch to CIFAR-10 — 32 × 32 colour images, ten classes — and run the identical script.
  • Replace the first Linear layer with nn.Conv2d. Convolutions exploit the fact that neighbouring pixels are related, and the accuracy jump is usually obvious.
  • Build a confusion matrix to see which classes get mixed up. Sneakers and ankle boots are a classic pair.
  • Add a learning rate schedule once your baseline is stable.

The first classifier is rarely the point. What you keep is the shape of the script: load data, define a model, pick a loss, loop, evaluate. Every project after this one is a variation on those five moves, and now you have made them once with your own hands.

Photo: Pexels / Pixabay

Copyright © 2026 Deep Learning Daily. All rights reserved. Deep Learning Daily.