thinc_mnist_mlp_training_with_relu_softmax_adam.py

1from thinc.api import chain, Relu, Softmax, Adam, fix_random_seed
2import ml_datasets
3
4# Set the random seed for reproducibility
5fix_random_seed(0)
6
7# 1. Load the MNIST dataset
8(train_X, train_Y), (test_X, test_Y) = ml_datasets.mnist()
9
10# 2. Define the model (a simple MLP: 784 -> 128 (ReLU) -> 10 (Softmax))
11model = chain(
12    Relu(nO=128, nI=784), 
13    Softmax(nO=10, nI=128)
14)
15
16# 3. Initialize the model with sample data to infer shapes
17model.initialize(X=train_X[:5], Y=train_Y[:5])
18
19# 4. Create the optimizer
20optimizer = Adam()
21
22# 5. Training loop
23for i in range(10):
24    # Get predictions and a callback to complete the backpropagation
25    yh, backprop = model.begin_update(train_X)
26    
27    # Calculate gradient of the loss (cross-entropy gradient is simple with Softmax)
28    dy = yh - train_Y
29    
30    # Backpropagate the gradient through the model
31    backprop(dy)
32    
33    # Update the model weights
34    model.finish_update(optimizer)
35    
36    # Evaluate the model
37    score = (model.predict(test_X).argmax(axis=1) == test_Y.argmax(axis=1)).mean()
38    print(f"Epoch {i}: accuracy {score:.3f}")