Classification with a network

Examples

For classification we change two things. The output layer has one unit per class, and the loss is the cross-entropy.

In torch the two are combined. CrossEntropyLoss expects the raw outputs, without a softmax, and applies the softmax internally. This is numerically more stable than doing it in two steps.

The network

import torch

model = torch.nn.Sequential(
    torch.nn.Linear(784, 128), torch.nn.ReLU(),
    torch.nn.Linear(128, 64), torch.nn.ReLU(),
    torch.nn.Linear(64, 10),          # 10 outputs, one per digit, no softmax
)

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

Loading MNIST

import numpy as np
from sklearn.datasets import fetch_openml

mnist = fetch_openml("mnist_784", version=1, as_frame=True, parser="auto")

X = mnist.data.values.astype(np.float32) / 255.0     # scale to [0, 1]
y = mnist.target.values.astype(np.int64)

X_train, y_train = X[:60000], y[:60000]
X_test, y_test = X[60000:], y[60000:]

Training

Xt = torch.tensor(X_train)
yt = torch.tensor(y_train)

loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(Xt, yt), batch_size=128, shuffle=True)

for epoch in range(10):
    for xb, yb in loader:
        opt.zero_grad()
        loss = loss_fn(model(xb), yb)     # raw outputs go in
        loss.backward()
        opt.step()
    print(epoch, round(float(loss), 4))

Evaluating

with torch.no_grad():
    logits = model(torch.tensor(X_test))
    pred = logits.argmax(dim=1).numpy()

print("test accuracy:", (pred == y_test).mean())

A network of this size reaches about 98 percent on MNIST. A multinomial logistic regression reaches about 92 percent. The difference is the learned features.

The confusion matrix

from sklearn.metrics import confusion_matrix
import pandas as pd

pd.DataFrame(confusion_matrix(y_test, pred))

Look at where the errors are. Fours confused with nines, threes with fives. These are the pairs that people also confuse, which is a sign that the network has learned something about shape rather than about pixel positions.

A small demonstration

The code above needs a download and a few minutes of training, so it is not run on this page. Here is a smaller version on the digits data set that ships with sklearn, which shows the same steps.

import numpy as np
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

torch.manual_seed(7)

digits = load_digits()
Xtr, Xte, ytr, yte = train_test_split(digits.data, digits.target,
                                      test_size=0.3, random_state=0)

sc = StandardScaler().fit(Xtr)
Xtr_t = torch.tensor(sc.transform(Xtr), dtype=torch.float32)
Xte_t = torch.tensor(sc.transform(Xte), dtype=torch.float32)
ytr_t = torch.tensor(ytr, dtype=torch.long)

model = torch.nn.Sequential(
    torch.nn.Linear(64, 64), torch.nn.ReLU(),
    torch.nn.Linear(64, 10),
)
loss_fn = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=1e-2)

loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(Xtr_t, ytr_t), batch_size=64, shuffle=True)

for epoch in range(30):
    for xb, yb in loader:
        opt.zero_grad()
        loss_fn(model(xb), yb).backward()
        opt.step()

with torch.no_grad():
    pred = model(Xte_t).argmax(dim=1).numpy()
print("test accuracy:", round(float((pred == yte).mean()), 4))

test accuracy: 0.9815

import matplotlib.pyplot as plt

wrong = np.where(pred != yte)[0][:8]
fig, axes = plt.subplots(1, len(wrong), figsize=(6.6, 1.3))
for ax, i in zip(axes, wrong):
    ax.imshow(Xte[i].reshape(8, 8), cmap="gray")
    ax.set_title(f"{yte[i]} → {pred[i]}", fontsize=8)
    ax.set(xticks=[], yticks=[])
plt.show()

Figure 37.1: Some wrongly classified test images, with the true label and the prediction.