Your first PyTorch model

Time to put the pieces together. In this lesson you will train a real classifier on Fashion-MNIST — 70,000 small grayscale images of clothing (T-shirts, trousers, sandals, bags) in 10 classes. It is small enough to train on a laptop CPU in a few minutes yet harder than handwritten digits, so it is a good first project for anyone interested in apparel or e-commerce catalogues.

Step 1: Load the data

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

tfm = transforms.ToTensor()   # PIL image -> float tensor in [0, 1], shape (1, 28, 28)
train_ds = datasets.FashionMNIST("data", train=True, download=True, transform=tfm)
test_ds = datasets.FashionMNIST("data", train=False, download=True, transform=tfm)

train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
test_dl = DataLoader(test_ds, batch_size=256)

x, y = next(iter(train_dl))
print(x.shape, y.shape)       # torch.Size([64, 1, 28, 28]) torch.Size([64])
print(train_ds.classes)

Training data is shuffled so each batch is a fair mix of classes; test data does not need shuffling.

Step 2: Define the model

Subclassing nn.Module is the standard pattern. Layers go in __init__, the computation goes in forward.

class FashionNet(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),              # (N, 1, 28, 28) -> (N, 784)
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, n_classes), # raw scores (logits)
        )

    def forward(self, x):
        return self.net(x)

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

Notice there is no softmax at the end. nn.CrossEntropyLoss expects raw logits and applies log-softmax internally, which is more numerically stable.

Step 3: Loss and optimizer

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Step 4: Train and evaluate

def train_one_epoch():
    model.train()
    total = 0.0
    for xb, yb in train_dl:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total += loss.item() * xb.size(0)
    return total / len(train_dl.dataset)

@torch.no_grad()
def evaluate():
    model.eval()
    correct = 0
    for xb, yb in test_dl:
        xb, yb = xb.to(device), yb.to(device)
        preds = model(xb).argmax(dim=1)
        correct += (preds == yb).sum().item()
    return correct / len(test_dl.dataset)

for epoch in range(5):
    loss = train_one_epoch()
    acc = evaluate()
    print(f"epoch {epoch + 1}: train loss {loss:.3f}, test acc {acc:.3%}")

Expect roughly 87–89% test accuracy after five epochs. That is a solid baseline; the CNN lesson will push it above 90%.

Step 5: Predict and save

model.eval()
img, label = test_ds[0]
with torch.no_grad():
    probs = torch.softmax(model(img.unsqueeze(0).to(device)), dim=1)
print("predicted:", train_ds.classes[probs.argmax().item()],
      "| actual:", train_ds.classes[label])

torch.save(model.state_dict(), "fashion_net.pt")

# later, in another script
restored = FashionNet()
restored.load_state_dict(torch.load("fashion_net.pt", map_location="cpu"))
restored.eval()

Save the state_dict (the weights), not the whole model object; it is portable and robust to code refactoring.

What each piece does

Piece	Role
Dataset	Knows how to fetch one example
DataLoader	Batches, shuffles and loads in parallel
nn.Module	Holds layers and parameters, defines forward
Loss function	Scores predictions against labels
Optimizer	Updates parameters from gradients
model.train() / model.eval()	Switch dropout and batch-norm behaviour
torch.no_grad()	Disables gradient tracking for faster inference

Common mistakes

• Adding nn.Softmax before CrossEntropyLoss — the loss then applies softmax twice and training stalls.
• Forgetting model.eval() during evaluation.
• Moving the model to GPU but not the batches (or vice versa).
• Printing only the last batch's loss, which is noisy; average over the epoch instead.