Mini-batch Training with DataLoader

Train a digit classifier the way real PyTorch code does: a Dataset, a DataLoader that shuffles and batches it, and a loop over epochs.

The data is scikit-learn's handwritten digits. Each image is 8 by 8 pixels, given as a row of 64 values from 0 to 16, and each label is a digit from 0 to 9.

Write two things.

1. DigitsDataset(X, y), a torch.utils.data.Dataset. It takes numpy arrays, keeps the pixels as float32 divided by 16 (so they are in [0, 1]) and the labels as int64. __len__ returns the number of samples. __getitem__(i) returns one pair (x_i, y_i): a tensor of shape (64,) and a scalar tensor.
2. train_digits(X_train, y_train), which wraps the dataset in a DataLoader with shuffle=True and a batch size such as 64, trains a network on mini-batches for some epochs, and returns the trained nn.Module. The module maps a batch of shape (B, 64) to logits of shape (B, 10).

The DataLoader calls __getitem__ once per sample and stacks the results into batches for you. That is why __getitem__ returns one sample, not a batch.

This is scored on the result. The tests build a DigitsDataset from held-out images and your model must reach at least 95% accuracy on it.

Function signature

class DigitsDataset(torch.utils.data.Dataset):
    def __init__(self, X, y): ...
    def __len__(self): ...
    def __getitem__(self, i): ...


def train_digits(X_train, y_train):

Example

Input

ds = DigitsDataset(X_train, y_train)
x0, y0 = ds[0]
model = train_digits(X_train, y_train)

Output

len(ds)              ->  1300
x0.shape, x0.dtype   ->  torch.Size([64]), torch.float32
y0.dtype             ->  torch.int64
model(x0[None]).shape  ->  torch.Size([1, 10])

Reasoning

The dataset holds 1300 samples and returns them one at a time. The model takes a batch, here a batch of one made with x0[None], and returns one row of 10 logits for each sample.

Constraints and notes

• Return the model in evaluation mode, with model.eval().
• nn.CrossEntropyLoss takes raw logits and int64 labels.
• One hidden layer of about 128 units, Adam with lr=1e-3 and about 30 epochs is enough. Set a seed with torch.manual_seed so the result is the same each run.
• Do not change the arrays you are given.