🎓 A1Learner
🔍

PyTorch Training the Classifier

A multi-class classifier is a model that picks one class from three or more. It gives one score per class for each row, and the highest score wins.

Labels for three classes

torch.full builds one label block per class. torch.cat joins the blocks in order.

labels = torch.cat([torch.full((3,), i, dtype=torch.long)
                    for i in range(3)])
print(labels)

An MLP with three outputs

Two inputs enter, and three scores leave. One score belongs to each class.

torch.manual_seed(1)
model = nn.Sequential(
    nn.Linear(2, 16),
    nn.ReLU(),
    nn.Linear(16, 3),
)
points = torch.randn(4, 2)
print(model(points).shape)

Batches from a DataLoader

The loader cuts the dataset into batches. The last batch holds the leftover rows.

dataset = TensorDataset(torch.randn(10, 2),
                        torch.zeros(10, dtype=torch.long))
loader = DataLoader(dataset, batch_size=4)
for batch_points, batch_labels in loader:
    print(batch_points.shape, batch_labels.shape)

Example

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)

per_class = 60
centres = [torch.tensor([-3.0, -2.0]), torch.tensor([3.0, -2.0]),
           torch.tensor([0.0, 3.0])]
blobs = [torch.randn(per_class, 2) * 0.8 + centre for centre in centres]
points = torch.cat(blobs)
labels = torch.cat([torch.full((per_class,), i, dtype=torch.long)
                     for i in range(3)])

dataset = TensorDataset(points, labels)
torch.manual_seed(2)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

torch.manual_seed(1)
model = nn.Sequential(
    nn.Linear(2, 16),
    nn.ReLU(),
    nn.Linear(16, 3),
)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(30):
    for batch_points, batch_labels in loader:
        scores = model(batch_points)
        loss = loss_fn(scores, batch_labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    predictions = model(points).argmax(dim=1)
    correct = int((predictions == labels).sum())
    accuracy = correct / labels.shape[0] * 100

print("Points:", points.shape[0])
print("Correct:", correct, "of", labels.shape[0])
print(f"Final accuracy: {accuracy:.0f}%")

Key points

  • Multi-class: one correct class per row, chosen from three or more
  • Use nn.Sequential to stack Linear, ReLU, and a final Linear
  • The last layer gives one output per class
  • CrossEntropyLoss takes the scores and the correct class numbers
  • Use a DataLoader to feed the model shuffled batches
  • Three scores per row carry the answer, and argmax reads the winner

In this example, three blobs of 60 points sit around three centres. The labels are class numbers 0, 1, and 2, in int64. A TensorDataset pairs the points with the labels, and a DataLoader serves batches of 32 with shuffling. A seed is set before the model is built. The model is an MLP with a hidden layer of 16 and three outputs. CrossEntropyLoss and Adam train it for 30 epochs. At the end, argmax reads one prediction per point, and all 180 predictions are right.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the final Linear layer of a three-class MLP output for each row?