PyTorch MNIST

PyTorch Hello, World! Very basic example of training, storing, and using a PyTorch model to see how it fits together.

Trains a simple MLP on the MNIST handwriting dataset from Hugging Face, saves the model. You can then run the model against a sample of the test data.

Extensions could be:

• add s3 store/load to the serialisers
• add setup/teardown infra to run the model on a GPU accelerated machine on AWS

This is not production quality code. It's not tested, secure, or performant!

Setup

uv sync

Usage

Train model and store weights

uv run main.py train

Run model

uv run main.py run                       # uses latest model in out/
uv run main.py run out/1744382400000.pt  # use a specific model

Dev

uv run ruff format .       # format                                                                                                                                                                                                                                                                              
uv run ruff check . --fix  # lint + autofix