A Small CNN for Digits

Build a convolutional network and train it to read handwritten digits.

The images are scikit-learn's digits: 8 by 8 pixels, one grey channel, already scaled to [0, 1]. They arrive in PyTorch's image layout, (N, C, H, W), which here is (N, 1, 8, 8). Labels are int64 digits from 0 to 9.

Write two things.

1. DigitCNN, an nn.Module that maps (N, 1, 8, 8) to logits of shape (N, 10). Use at least one nn.Conv2d. A good shape is two blocks of Conv2d(3x3, padding=1) -> ReLU -> MaxPool2d(2), which takes the image from 8x8 to 4x4 to 2x2, then flatten and one nn.Linear.
2. train_cnn(X_train, y_train), which trains a DigitCNN on mini-batches and returns it in evaluation mode.

This is scored on the result: the returned model must reach at least 95% accuracy on held-out images.

Function signature

class DigitCNN(nn.Module):
    def __init__(self): ...
    def forward(self, x): ...


def train_cnn(X_train, y_train):

Example

Input

net = DigitCNN()
net(torch.zeros(5, 1, 8, 8)).shape

Output

torch.Size([5, 10])

Reasoning

With 16 filters, then 32, and two poolings, the feature map is (5, 32, 2, 2). flatten(1) keeps the batch axis and gives (5, 128). A Linear(128, 10) gives one row of 10 logits for each image.

Constraints and notes

• Flatten with x.flatten(1), not x.view(-1) or x.squeeze(). The model must work for a batch of one image.
• Return the model after model.eval().
• Adam with lr=3e-3, batch size 64 and about 25 epochs is enough. Set a seed with torch.manual_seed.
• Keep training short. The grader runs on a CPU.