pytorch-mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1353ad3a",
   "metadata": {},
   "source": [
    "# PyTorch 3LC MNIST Sample Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ddc779e",
   "metadata": {},
   "source": [
    "<div style=\"display: inline-flex; align-items: center; gap: 10px;\">\n",
    "        <a href=\"https://colab.research.google.com/github/3lc-ai/notebook-examples/blob/main/pytorch-mnist.ipynb\"\n",
    "        target=\"_blank\"\n",
    "            style=\"background-color: transparent; text-decoration: none; display: inline-flex; align-items: center;\n",
    "            padding: 5px 10px; font-family: Arial, sans-serif;\"> <img\n",
    "            src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" style=\"height: 30px;\n",
    "            vertical-align: middle;box-shadow: none;\"/>\n",
    "        </a> <a href=\"https://github.com/3lc-ai/notebook-examples/blob/main/pytorch-mnist.ipynb\"\n",
    "            style=\"text-decoration: none; display: inline-flex; align-items: center; background-color: #ffffff; border:\n",
    "            1px solid #d1d5da; border-radius: 8px; padding: 2px 10px; color: #333; font-family: Arial, sans-serif;\">\n",
    "            <svg aria-hidden=\"true\" focusable=\"false\" role=\"img\" class=\"octicon octicon-mark-github\" viewBox=\"0 0 16 16\"\n",
    "            width=\"20\" height=\"20\" fill=\"#333\"\n",
    "            style=\"display:inline-block;user-select:none;vertical-align:text-bottom;overflow:visible; margin-right:\n",
    "            8px;\">\n",
    "                <path d=\"M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2\n",
    "                0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0\n",
    "                0-.67-.22-2.2.82-.64-.18-1.32-.27-2-.27-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16\n",
    "                1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51\n",
    "                1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68\n",
    "                1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z\"></path>\n",
    "            </svg> <span style=\"vertical-align: middle; color: #333;\">Open in GitHub</span>\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a0c11b9f",
   "metadata": {
    "papermill": {
     "duration": 0.008129,
     "end_time": "2023-04-19T12:15:44.827629",
     "exception": false,
     "start_time": "2023-04-19T12:15:44.819500",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "This notebook demonstrates the training of a Convolutional Neural Network (CNN) on the MNIST dataset using PyTorch and\n",
    "3LC. The built-in MNIST dataset for training and validation is wrapped in a Table. The training runs for 5\n",
    "epochs, and during this period, classification metrics and embeddings are collected.\n",
    "\n",
    "The notebook demonstrates:\n",
    "\n",
    "+ How to use a 3LC Table for integrating with built-in PyTorch datasets.\n",
    "+ Metrics collection using a custom `MetricsCollector` subclass and a `EmbeddingsMetricsCollector`.\n",
    "+ Reducing the dimensionality of embeddings using PaCMAP as a post-processing step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb0fd44",
   "metadata": {},
   "source": [
    "## Project Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c484e1d9",
   "metadata": {
    "papermill": {
     "duration": 0.034774,
     "end_time": "2023-04-19T12:15:44.870371",
     "exception": false,
     "start_time": "2023-04-19T12:15:44.835597",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "PROJECT_NAME = \"MNIST Digit Classification\"\n",
    "RUN_NAME = \"Train MNIST Classifier\"\n",
    "DESCRIPTION = \"Train a simple CNN to classify MNIST digits\"\n",
    "TRAIN_DATASET_NAME = \"mnist-train\"\n",
    "VAL_DATASET_NAME = \"mnist-val\"\n",
    "TRANSIENT_DATA_PATH = \"../transient_data\"\n",
    "COLLECT_METRICS_BATCH_SIZE = 2048\n",
    "TRAIN_BATCH_SIZE = 64\n",
    "INITIAL_LR = 1.0\n",
    "LR_GAMMA = 0.7\n",
    "EPOCHS = 5\n",
    "NUM_WORKERS = 0\n",
    "DEVICE = \"cuda:0\"\n",
    "TLC_PUBLIC_EXAMPLES_DEVELOPER_MODE = True\n",
    "INSTALL_DEPENDENCIES = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09db29b4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "if INSTALL_DEPENDENCIES:\n",
    "    %pip --quiet install torch --index-url https://download.pytorch.org/whl/cu118\n",
    "    %pip --quiet install torchvision --index-url https://download.pytorch.org/whl/cu118\n",
    "    %pip --quiet install 3lc[pacmap]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c9de34",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b59fda27",
   "metadata": {
    "papermill": {
     "duration": 1.254606,
     "end_time": "2023-04-19T12:15:46.131975",
     "exception": false,
     "start_time": "2023-04-19T12:15:44.877369",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import tlc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61ecf72b",
   "metadata": {},
   "source": [
    "## Initialize a 3LC Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5119e73c",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"train_batch_size\": TRAIN_BATCH_SIZE,\n",
    "    \"initial_lr\": INITIAL_LR,\n",
    "    \"lr_gamma\": LR_GAMMA,\n",
    "    \"epochs\": EPOCHS,\n",
    "}\n",
    "\n",
    "run = tlc.init(\n",
    "    project_name=PROJECT_NAME,\n",
    "    run_name=RUN_NAME,\n",
    "    description=DESCRIPTION,\n",
    "    parameters=config,\n",
    "    if_exists=\"overwrite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc68a489",
   "metadata": {},
   "source": [
    "## Setup Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "482397d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = torchvision.transforms.Compose(\n",
    "    [torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.1307,), (0.3081,))]\n",
    ")\n",
    "\n",
    "train_dataset = torchvision.datasets.MNIST(root=TRANSIENT_DATA_PATH, train=True, download=True)\n",
    "eval_dataset = torchvision.datasets.MNIST(root=TRANSIENT_DATA_PATH, train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4574d138",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = [str(i) for i in range(10)]\n",
    "\n",
    "structure = (tlc.PILImage(\"image\"), tlc.CategoricalLabel(\"label\", class_names))\n",
    "\n",
    "\n",
    "def transforms(x):\n",
    "    return transform(x[0]), torch.tensor(x[1])\n",
    "\n",
    "\n",
    "# We pick up the latest version of the dataset, so that we can re-run this notebook as-is\n",
    "# after adding new revisions to the dataset.\n",
    "tlc_train_dataset = (\n",
    "    tlc.Table.from_torch_dataset(\n",
    "        dataset=train_dataset,\n",
    "        dataset_name=TRAIN_DATASET_NAME,\n",
    "        structure=structure,\n",
    "        project_name=PROJECT_NAME,\n",
    "        description=\"MNIST training dataset\",\n",
    "        table_name=\"train\",\n",
    "        if_exists=\"overwrite\",\n",
    "    )\n",
    "    .map(transforms)\n",
    "    .latest()\n",
    ")\n",
    "\n",
    "tlc_val_dataset = (\n",
    "    tlc.Table.from_torch_dataset(\n",
    "        dataset=eval_dataset,\n",
    "        dataset_name=VAL_DATASET_NAME,\n",
    "        structure=structure,\n",
    "        project_name=PROJECT_NAME,\n",
    "        description=\"MNIST validation dataset\",\n",
    "        table_name=\"val\",\n",
    "        if_exists=\"overwrite\",\n",
    "    )\n",
    "    .map(transforms)\n",
    "    .latest()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cc2a06b",
   "metadata": {},
   "source": [
    "## Setup Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54424d08",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = torch.device(DEVICE)\n",
    "print(f\"Using device: {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c943fba2",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    # From https://github.com/pytorch/examples/blob/main/mnist/main.py\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "\n",
    "model = Net().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d918ed7a",
   "metadata": {},
   "source": [
    "## Setup Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a021b68",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adadelta(model.parameters(), lr=INITIAL_LR)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=LR_GAMMA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6faa8d8",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for data, target in tqdm(train_loader, desc=f\"Training {epoch+1}/{EPOCHS}\"):  # Epoch is 0-indexed\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e569b26c",
   "metadata": {},
   "source": [
    "## Setup Metrics Collectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c18f87",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTMetricsCollector(tlc.MetricsCollector):\n",
    "    def __init__(self, criterion):\n",
    "        super().__init__()\n",
    "        self._criterion = criterion\n",
    "\n",
    "    def compute_metrics(self, batch, predictor_output):\n",
    "        predictions = predictor_output.forward\n",
    "        labels = batch[1].to(device)\n",
    "\n",
    "        metrics = {\n",
    "            \"loss\": self._criterion(predictions, labels).cpu().numpy(),\n",
    "            \"predicted\": torch.argmax(predictions, dim=1).cpu().numpy(),\n",
    "            \"confidence\": torch.exp(torch.max(predictions, dim=1).values).cpu().numpy(),\n",
    "            \"accuracy\": (torch.argmax(predictions, dim=1) == labels).cpu().numpy(),\n",
    "        }\n",
    "        return metrics\n",
    "\n",
    "    @property\n",
    "    def column_schemas(self):\n",
    "        # Explicitly override the schema of the predicted label, in order for it to be displayed as a\n",
    "        # categorical label in the Dashboard.\n",
    "        schemas = {\n",
    "            \"predicted\": tlc.CategoricalLabelSchema(\n",
    "                class_names,\n",
    "                display_name=\"predicted label\",\n",
    "            )\n",
    "        }\n",
    "        return schemas\n",
    "\n",
    "\n",
    "mnist_metrics_collector = MNISTMetricsCollector(nn.NLLLoss(reduction=\"none\"))\n",
    "embeddings_metrics_collector = tlc.EmbeddingsMetricsCollector(layers=[4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "118ff375",
   "metadata": {},
   "source": [
    "## Run Training\n",
    "\n",
    "We run training using a weighted sampler provided by the 3LC Table. The sampler uses the default `weights` column\n",
    "to sample the data. The weights can be updated in the Dashboard, and will be automatically picked up by the sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0554e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = tlc_train_dataset.create_sampler()\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    tlc_train_dataset,\n",
    "    batch_size=TRAIN_BATCH_SIZE,\n",
    "    sampler=sampler,\n",
    "    num_workers=NUM_WORKERS,\n",
    ")\n",
    "\n",
    "metrics_collection_dataloader_args = {\n",
    "    \"num_workers\": NUM_WORKERS,\n",
    "    \"batch_size\": COLLECT_METRICS_BATCH_SIZE,\n",
    "}\n",
    "\n",
    "predictor = tlc.Predictor(model, layers=[4])\n",
    "\n",
    "# Train the model\n",
    "for epoch in range(EPOCHS):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "\n",
    "    tlc.collect_metrics(\n",
    "        tlc_train_dataset,\n",
    "        metrics_collectors=[\n",
    "            mnist_metrics_collector,\n",
    "            embeddings_metrics_collector,\n",
    "        ],\n",
    "        predictor=predictor,\n",
    "        split=\"train\",\n",
    "        constants={\"epoch\": epoch},\n",
    "        dataloader_args=metrics_collection_dataloader_args,\n",
    "    )\n",
    "    tlc.collect_metrics(\n",
    "        tlc_val_dataset,\n",
    "        metrics_collectors=[\n",
    "            mnist_metrics_collector,\n",
    "            embeddings_metrics_collector,\n",
    "        ],\n",
    "        predictor=predictor,\n",
    "        split=\"val\",\n",
    "        constants={\"epoch\": epoch},\n",
    "        dataloader_args=metrics_collection_dataloader_args,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a8577d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reduce embeddings using the final validation-set embeddings to fit a PaCMAP model\n",
    "url_mapping = run.reduce_embeddings_by_foreign_table_url(\n",
    "    tlc_val_dataset.url,\n",
    "    method=\"pacmap\",\n",
    "    n_components=3,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.8.10"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 5.341171,
   "end_time": "2023-04-19T12:15:47.738497",
   "environment_variables": {},
   "exception": true,
   "input_path": "c:\\Users\\RupalDangi\\experimental-notebooks\\experimental-notebooks\\vision\\classification\\MNIST\\torch-examples-cnn-tlc.ipynb",
   "output_path": "c:\\Users\\RupalDangi\\experimental-notebooks\\experimental-notebooks\\vision\\classification\\MNIST\\torch-examples-cnn-tlc.ipynb",
   "parameters": {},
   "start_time": "2023-04-19T12:15:42.397326",
   "version": "2.4.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "4ee08d3161a7c3d0b4fb68735cf1133b294ca2a41e3afd68fe473bc7561f5f08"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}