{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Train a Model on MNIST with FiftyOne and Torch\n",
    "\n",
    "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:\n",
    "\n",
    "* Loading the MNIST dataset from the [Dataset Zoo](https://docs.voxel51.com/user_guide/dataset_zoo/index.html)\n",
    "* Creating train/validation/test splits with FiftyOne's tagging and random splitting utilities\n",
    "* Building a subset of the dataset for faster experimentation\n",
    "* Running a simple training loop via an external script ([mnist_training.py](https://github.com/voxel51/fiftyone/blob/develop/docs/source/recipes/torch-dataset-examples/mnist_training.py))\n",
    "* Saving model weights for later evaluation or reuse\n",
    "\n",
    "**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)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "If you haven't already, install FiftyOne:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install fiftyone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fiftyone as fo\n",
    "import fiftyone.zoo as foz\n",
    "import fiftyone.utils.random as four"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "import torchvision.transforms.v2 as transforms\n",
    "from torchvision import tv_tensors\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as plt_patches\n",
    "from PIL import Image\n",
    "import urllib.request"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run this recipe, you'll need the `mnist_training.py` script ([source on GitHub](https://github.com/voxel51/fiftyone/blob/develop/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.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://cdn.voxel51.com/tutorials_torch_dataset_examples/notebook_simple_training_example/mnist_training.py\"\n",
    "urllib.request.urlretrieve(url, \"mnist_training.py\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# utils.py is shared across the torch-dataset-examples notebooks\n",
    "url = \"https://cdn.voxel51.com/tutorials_torch_dataset_examples/notebook_the_cache_field_names_argument/utils.py\"\n",
    "urllib.request.urlretrieve(url, \"utils.py\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mnist_training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.multiprocessing.set_start_method(\"forkserver\")\n",
    "torch.multiprocessing.set_forkserver_preload([\"torch\", \"fiftyone\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training MNIST with FiftyOneTorchDataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist = foz.load_zoo_dataset(\"mnist\")\n",
    "mnist.persistent = True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "Before defining splits, you can optionally explore the dataset in the FiftyOne App:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fo.launch_app(mnist, auto=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define a validation split from the non-test samples. FiftyOne's `random_split` makes this straightforward:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove existing 'train' or 'validation' tags if they exist\n",
    "mnist.untag_samples([\"train\", \"validation\"])\n",
    "\n",
    "# create a random split on all non-test samples\n",
    "not_test = mnist.match_tags(\"test\", bool=False)\n",
    "four.random_split(not_test, {\"train\": 0.9, \"validation\": 0.1})\n",
    "print(mnist.count_sample_tags())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a small subset for faster experimentation\n",
    "samples = []\n",
    "samples += mnist.match_tags(\"train\").take(1000).values(\"id\")\n",
    "for tag in [\"test\", \"validation\"]:\n",
    "    samples += mnist.match_tags(tag).values(\"id\")\n",
    "\n",
    "subset = mnist.select(samples)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "path_to_save_weights = Path(\"./mnist_weights\")\n",
    "path_to_save_weights.mkdir(parents=True, exist_ok=True)\n",
    "mnist_training.main(subset, 10, 10, device, str(path_to_save_weights))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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:\n",
    "\n",
    "```python\n",
    "fo.launch_app(mnist)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "## Understanding the Training Script\n",
    "\n",
    "The [mnist_training.py](https://github.com/voxel51/fiftyone/blob/develop/docs/source/recipes/torch-dataset-examples/mnist_training.py) script contains the full training loop. Here we walk through its key design decisions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "### DataLoader Creation with FiftyOneTorchDataset and GetItem\n",
    "\n",
    "`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:\n",
    "\n",
    "```python\n",
    "class MnistGetItem(GetItem):\n",
    "    def __init__(self):\n",
    "        super().__init__(\n",
    "            field_mapping={\"id\": \"id\", \"filepath\": \"filepath\", \"label\": \"ground_truth.label\"}\n",
    "        )\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image = convert_and_normalize(Image.open(sample[\"filepath\"]).convert(\"RGB\"))\n",
    "        label = int(sample[\"label\"][0])\n",
    "        return {\"image\": image, \"label\": label, \"id\": sample[\"id\"]}\n",
    "```\n",
    "\n",
    "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.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "### Versatility: Any View Becomes a Split\n",
    "\n",
    "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:\n",
    "\n",
    "```python\n",
    "# 90/10 random split from all non-test samples\n",
    "not_test = mnist.match_tags(\"test\", bool=False)\n",
    "four.random_split(not_test, {\"train\": 0.9, \"validation\": 0.1})\n",
    "\n",
    "# Or scope training to a curated subset\n",
    "subset = mnist.select(selected_ids)\n",
    "```\n",
    "\n",
    "You can pass any view — filtered by tag, label, quality metric, or brain-run result — straight into `create_dataloaders()` without changing the training script.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "### Writing Predictions Back to FiftyOne\n",
    "\n",
    "During evaluation, the script writes per-sample predictions back to the dataset:\n",
    "\n",
    "```python\n",
    "fo_predictions = [\n",
    "    fo.Classification(\n",
    "        label=utils.mnist_index_to_label_string(np.argmax(sample_logits)),\n",
    "        logits=sample_logits,\n",
    "    )\n",
    "    for sample_logits in prediction.detach().cpu().numpy()\n",
    "]\n",
    "samples.set_values(\"predictions\", fo_predictions)\n",
    "samples.save()\n",
    "```\n",
    "\n",
    "After training, you can open the FiftyOne App and immediately browse predictions, filter by confidence, and inspect misclassified samples — all within the same workflow.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "",
   "metadata": {},
   "source": [
    "### Evaluation with FiftyOne\n",
    "\n",
    "Once predictions are stored, FiftyOne's built-in evaluation API runs directly on the test split:\n",
    "\n",
    "```python\n",
    "results = dataset.match_tags(\"test\").evaluate_classifications(\n",
    "    \"predictions\",\n",
    "    gt_field=\"ground_truth\",\n",
    "    eval_key=\"eval\",\n",
    "    classes=classes,\n",
    "    k=3,\n",
    ")\n",
    "results.print_report(classes=classes)\n",
    "```\n",
    "\n",
    "Results are persisted under the `eval` key, making per-class metrics and top-k accuracy available for review in the App at any time.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch-dataset",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}