Choosing Among Many: Classification with Softmax

# Chapter 9: more than two classes. Softmax turns a layer's scores into probabilities.
# Uses the Tensor class from Chapter 6, saved as tensor.py (see Chapter 7).
import numpy as np
from tensor import Tensor

ACTIVITIES = ["picnic", "kite", "read"]


def softmax(z):
    """Plain NumPy, one row of scores per example. Subtracting the row's largest score changes nothing but avoids overflow."""
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


# === 1. Chapter 2's layer, scored on tomorrow's weather ===========================
W = np.array([[2.0, -3.0, -1.0],     # picnic
              [-1.0, 0.0, 3.0],      # kite flying
              [0.0, 2.5, 0.0]])      # stay in and read
b = np.array([0.5, -0.5, 0.0])
tomorrow = np.array([0.7, 0.2, 0.5])                 # (sun, rain, wind)
z = W @ tomorrow + b
print("scores z:", z, " e^z:", np.exp(z).round(4), " sum:", np.exp(z).sum().round(4))
print("softmax: ", softmax(z).round(4))
for T in [0.5, 2.0]:
    print(f"T = {T}:   ", softmax(z / T).round(4))

# Two classes is the sigmoid in disguise: softmax([z1, z2])[0] = sigmoid(z1 - z2)
print("picnic vs kite only:", softmax(z[:2]).round(4), " sigmoid(0.8 - 0.3) =", round(1 / (1 + np.exp(-0.5)), 4))

# Why we subtract the maximum first
big = np.array([1000.0, 999.0, 998.0])
with np.errstate(over="ignore", invalid="ignore"):
    naive = np.exp(big) / np.exp(big).sum()
print("naive softmax of [1000, 999, 998]:", naive, "  shifted:", softmax(big).round(4))

# === 2. The loss, and its gradient p − y ============================================
p = softmax(z)
y = np.array([0.0, 1.0, 0.0])                        # tomorrow they went kite flying
print(f"\ncross-entropy if they flew a kite: -ln {p[1]:.4f} = {-np.log(p[1]):.4f}   (uniform guess: ln 3 = {np.log(3):.4f})")
print("gradient with respect to z, p − y:", (p - y).round(4), " sums to", round(float((p - y).sum()), 12))


def softmax_cross_entropy(Z, Y):
    """Mean cross-entropy of softmax(Z) against one-hot rows Y, built from Chapter 6's operations only."""
    shifted = Z + (-Z.data.max(axis=1, keepdims=True))            # a constant shift: no gradient needed
    E = shifted.exp()
    row_sums = E @ Tensor(np.ones((Z.shape[1], 1)))               # (N, 3)(3, 1): each row's total
    P = E / row_sums                                              # broadcast over the 3 columns
    return -(Tensor(Y) * P.log()).sum() / Z.shape[0], P

# Check the engine against p − y: make the scores a leaf and ask for their gradient
Z = Tensor(z[None, :])
loss, P = softmax_cross_entropy(Z, y[None, :])
loss.backward()
print("engine's dL/dz:", Z.grad.round(4), " loss", round(float(loss.data), 4))

# === 3. A season of days, and a classifier trained on them ===========================
rng = np.random.default_rng(9)
X = np.round(rng.random((150, 3)), 2)                # (sun, rain, wind) for 150 days
rule = np.where(X[:, 1] > 0.5, 2, np.where(X[:, 2] > 0.6, 1, 0))   # rainy: read; else windy: kite; else picnic
labels = rule.copy()
noisy = rng.random(150) < 0.1                        # 1 day in 10, people did something else
labels[noisy] = rng.integers(0, 3, noisy.sum())
Y = np.eye(3)[labels]                                # one-hot rows, (150, 3)
print(f"\n150 days: {np.bincount(labels).tolist()} picnic/kite/read;"
      f" the rule itself is right on {np.mean(rule == labels):.0%} of them")


def train(params, scores, lr, steps):
    for _ in range(steps):
        loss, _ = softmax_cross_entropy(scores(), Y)
        for prm in params:
            prm.grad = np.zeros_like(prm.data)
        loss.backward()
        for prm in params:
            prm.data = prm.data - lr * prm.grad
    loss, P = softmax_cross_entropy(scores(), Y)
    return float(loss.data), np.mean(P.data.argmax(axis=1) == labels)

# Softmax regression: one linear layer, 3 inputs -> 3 scores, starting from all zeros
W = Tensor(np.zeros((3, 3)))
b = Tensor(np.zeros(3))
start, _ = softmax_cross_entropy(Tensor(X) @ W.T + b, Y)
loss, acc = train([W, b], lambda: Tensor(X) @ W.T + b, lr=1.0, steps=3000)
print(f"softmax regression: loss {float(start.data):.4f} -> {loss:.4f}, accuracy {acc:.1%}")
for name, row, bias in zip(ACTIVITIES, W.data, b.data):
    print(f"  {name:6} weights (sun, rain, wind) = {row.round(2)}, bias {bias:+.2f}")

# With a hidden layer of 8 or 16 ReLU units (Chapter 8)
for hidden in [8, 16]:
    r = np.random.default_rng(0)
    W1, b1 = Tensor(r.normal(size=(hidden, 3))), Tensor(np.zeros(hidden))
    W2, b2 = Tensor(r.normal(size=(3, hidden))), Tensor(np.zeros(3))
    loss, acc = train([W1, b1, W2, b2], lambda: (Tensor(X) @ W1.T + b1).relu() @ W2.T + b2, lr=0.5, steps=5000)
    print(f"{hidden:2} hidden units:     loss {loss:.4f}, accuracy {acc:.1%}")