# How to Train a Model on MNIST with FiftyOne and Torch This recipe demonstrates how to train a PyTorch model on the **MNIST** dataset using [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset). This is useful when you want to build and evaluate models in Torch while managing your data pipeline directly from FiftyOne. Specifically, it covers: - Loading the MNIST dataset from the [Dataset Zoo](https://docs.voxel51.com/user_guide/dataset_zoo/index.html) - Creating train/validation/test splits with FiftyOne’s tagging and random splitting utilities - Building a subset of the dataset for faster experimentation - Running a simple training loop via an external script ([mnist_training.py](https://github.com/voxel51/fiftyone/blob/main/docs/source/recipes/torch-dataset-examples/mnist_training.py)) - Saving model weights for later evaluation or reuse **API references:** [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset) · [GetItem](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.GetItem) ## Setup If you haven’t already, install FiftyOne: In this tutorial, we’ll use [PyTorch](https://pytorch.org/) for working with tensors and inspecting sample data. To follow along, you’ll need to install `torch` and `torchvision`, if necessary: ## Import Libraries To run this recipe, you’ll need the `mnist_training.py` script ([source on GitHub](https://github.com/voxel51/fiftyone/blob/main/docs/source/recipes/torch-dataset-examples/mnist_training.py)), which contains a complete PyTorch training loop built on [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset). The following cell downloads it into your working directory so it can be imported directly. ## Training MNIST with FiftyOneTorchDataset With our dataset loaded and splits defined, we can call `mnist_training.main()` directly. Under the hood this uses a [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset) and a [GetItem](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.GetItem) to build each split’s dataloader from FiftyOne views. Before defining splits, you can optionally explore the dataset in the FiftyOne App: Now let’s define a validation split from the non-test samples. FiftyOne’s `random_split` makes this straightforward: Training is complete. Predictions are written back to the underlying MNIST dataset via sample IDs, so opening `mnist` (the full dataset) will show all test-split predictions for review: ```python fo.launch_app(mnist) ``` ## Understanding the Training Script The [mnist_training.py](https://github.com/voxel51/fiftyone/blob/main/docs/source/recipes/torch-dataset-examples/mnist_training.py) script contains the full training loop. Here we walk through its key design decisions. ### DataLoader Creation with FiftyOneTorchDataset and GetItem `create_dataloaders()` calls `dataset.match_tags(split).to_torch(get_item)` for each split tag, converting every FiftyOne view directly into a [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset). The `MnistGetItem` class — a subclass of [GetItem](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.GetItem) — declares which fields to load and converts each sample into model-ready tensors: ```python class MnistGetItem(GetItem): def __init__(self): super().__init__( field_mapping={"id": "id", "filepath": "filepath", "label": "ground_truth.label"} ) def __call__(self, sample): image = convert_and_normalize(Image.open(sample["filepath"]).convert("RGB")) label = int(sample["label"][0]) return {"image": image, "label": label, "id": sample["id"]} ``` The resulting dataset plugs directly into `torch.utils.data.DataLoader`. The only required addition is `worker_init_fn=FiftyOneTorchDataset.worker_init`, which lets FiftyOne open its own database connection inside each worker process. ### Versatility: Any View Becomes a Split Because [FiftyOneTorchDataset](https://docs.voxel51.com/api/fiftyone.utils.torch.html#fiftyone.utils.torch.FiftyOneTorchDataset) is built from a FiftyOne *view*, any filtering, sorting, or tagging operation in FiftyOne automatically becomes a training or validation split — no data duplication needed. The cells above illustrate this: ```python # 90/10 random split from all non-test samples not_test = mnist.match_tags("test", bool=False) four.random_split(not_test, {"train": 0.9, "validation": 0.1}) # Or scope training to a curated subset subset = mnist.select(selected_ids) ``` You can pass any view — filtered by tag, label, quality metric, or brain-run result — straight into `create_dataloaders()` without changing the training script. ### Writing Predictions Back to FiftyOne During evaluation, the script writes per-sample predictions back to the dataset: ```python fo_predictions = [ fo.Classification( label=utils.mnist_index_to_label_string(np.argmax(sample_logits)), logits=sample_logits, ) for sample_logits in prediction.detach().cpu().numpy() ] samples.set_values("predictions", fo_predictions) samples.save() ``` After training, you can open the FiftyOne App and immediately browse predictions, filter by confidence, and inspect misclassified samples — all within the same workflow. ### Evaluation with FiftyOne Once predictions are stored, FiftyOne’s built-in evaluation API runs directly on the test split: ```python results = dataset.match_tags("test").evaluate_classifications( "predictions", gt_field="ground_truth", eval_key="eval", classes=classes, k=3, ) results.print_report(classes=classes) ``` Results are persisted under the `eval` key, making per-class metrics and top-k accuracy available for review in the App at any time.