Softmax Regression on Digits

Train multinomial logistic regression, also called softmax regression, to read handwritten digits.

The data is scikit-learn's load_digits: 8 by 8 images flattened to 64 pixel values from 0 to 16, with labels 0 to 9. Softmax regression is one linear layer that gives a score, a logit, for each class:

logits = X @ W.T + b              # (n, n_classes)
p      = softmax(logits)          # one probability per class
loss   = mean(-log p[i, y_i])     # cross-entropy

Train it with mini-batches. Wrap the tensors in a TensorDataset, use a DataLoader with shuffle=True, and make several passes (epochs) over the data. Return the trained module. It must map an (m, 64) tensor of raw pixel values to (m, n_classes) logits.

The model must be softmax regression, not a neural network. It must have exactly 64 * n_classes + n_classes trainable values. You can also store other tensors in it, for example the mean and standard deviation of the pixels, if you register them as buffers.

This problem is scored on the result. On held-out digits, your model must reach at least 94% accuracy.

Function signature

def train_softmax_regression(X_train, y_train, n_classes):

Example

Input

model = train_softmax_regression(X_train, y_train, n_classes=10)
model(X_val[:3])

Output

tensor([[ -3.1,  0.4, 11.8, ...],     # shape (3, 10)
        ...])

Reasoning

Each row holds ten logits. The largest one is the predicted digit, so the first image here is read as a 2. argmax of the logits gives the same class as argmax of the softmax probabilities.

Constraints and notes

• X_train is float32 with shape (n, 64). y_train is int64 with shape (n,).
• nn.CrossEntropyLoss takes raw logits. Do not put a softmax before it.
• Some pixels are always 0 in the training set. If you standardise, do not divide by a standard deviation of 0.
• Keep training small. A few hundred mini-batch steps are enough.