how-to-guides/hpo/notebooks/Neptune_HPO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to use Neptune Scale for tracking HPO runs\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/neptune-ai/scale-examples/blob/main/how-to-guides/hpo/notebooks/Neptune_HPO.ipynb\"> \n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"/> \n",
    "</a>\n",
    "<a target=\"_blank\" href=\"https://github.com/neptune-ai/scale-examples/blob/main/how-to-guides/hpo/notebooks/Neptune_HPO.ipynb\">\n",
    "  <img alt=\"Open in GitHub\" src=\"https://img.shields.io/badge/Open_in_GitHub-blue?logo=github&labelColor=black\">\n",
    "</a>\n",
    "<a target=\"_blank\" href=\"https://scale.neptune.ai/o/examples/org/hpo/runs/table?viewId=9d44261f-32a1-42e7-96ff-9b35edc4be66\">\n",
    "  <img alt=\"Explore in Neptune\" src=\"https://neptune.ai/wp-content/uploads/2024/01/neptune-badge.svg\">\n",
    "</a>\n",
    "<a target=\"_blank\" href=\"https://docs-beta.neptune.ai/tutorials/hpo/\">\n",
    "  <img alt=\"View tutorial in docs\" src=\"https://neptune.ai/wp-content/uploads/2024/01/docs-badge-2.svg\">\n",
    "</a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "When running a hyperparameter optimization job, you can use Neptune Scale to track all the metadata from the study and each trial.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before you start\n",
    "\n",
    "  1. Create a Neptune Scale account. [Register &rarr;](https://neptune.ai/early-access)\n",
    "  2. Create a Neptune project that you will use for tracking metadata. For instructions, see [Projects](https://docs-beta.neptune.ai/projects/) in the Neptune Scale docs.\n",
    "  3. Install and configure Neptune Scale for logging metadata. For instructions, see [Get started](https://docs-beta.neptune.ai/setup) in the Neptune Scale docs."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Neptune Scale and dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qU neptune-scale torch torchvision tqdm \"numpy<2.0\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neptune_scale import Run\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from tqdm.auto import trange, tqdm\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "ALLOWED_DATATYPES = [int, float, str, datetime, bool, list, set]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"batch_size\": 128,\n",
    "    \"input_size\": (1, 28, 28),\n",
    "    \"n_classes\": 10,\n",
    "    \"epochs\": 3,\n",
    "    \"device\": torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\"),\n",
    "}\n",
    "\n",
    "input_size = math.prod(parameters[\"input_size\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameter search space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rates = [0.025, 0.05, 0.075]  # learning rate choices"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseModel(nn.Module):\n",
    "    def __init__(self, input_size, num_classes):\n",
    "        super(BaseModel, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, 128)\n",
    "        self.fc2 = nn.Linear(128, 64)\n",
    "        self.fc3 = nn.Linear(64, num_classes)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_tfms = {\n",
    "    \"train\": transforms.Compose(\n",
    "        [\n",
    "            transforms.ToTensor(),\n",
    "        ]\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = datasets.MNIST(\n",
    "    root=\"mnist\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=data_tfms[\"train\"],\n",
    ")\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(\n",
    "    trainset,\n",
    "    batch_size=parameters[\"batch_size\"],\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log metadata across HPO trials into a single run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BaseModel(\n",
    "    input_size,\n",
    "    parameters[\"n_classes\"],\n",
    ").to(parameters[\"device\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a global Neptune run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "\n",
    "run = Run(run_id=f\"hpo-{uuid4()}\")\n",
    "\n",
    "run.add_tags([\"all-trials\", \"notebook\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Log configuration common across all trials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in parameters:\n",
    "    if type(parameters[key]) not in ALLOWED_DATATYPES:\n",
    "        run.log_configs({f\"config/{key}\": str(parameters[key])})\n",
    "    else:\n",
    "        run.log_configs({f\"config/{key}\": parameters[key]})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for trial, lr in tqdm(\n",
    "    enumerate(learning_rates),\n",
    "    total=len(learning_rates),\n",
    "    desc=\"Trials\",\n",
    "):\n",
    "    # Log trial hyperparameters\n",
    "    run.log_configs({f\"trials/{trial}/parameters/lr\": lr})\n",
    "\n",
    "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "    # Initialize fields for best values across all trials\n",
    "    best_acc = None\n",
    "\n",
    "    step = 0\n",
    "\n",
    "    for epoch in trange(parameters[\"epochs\"], desc=f\"Trial {trial} - lr: {lr}\"):\n",
    "        run.log_metrics(data={f\"trials/{trial}/epochs\": epoch}, step=epoch)\n",
    "\n",
    "        for x, y in trainloader:\n",
    "            x, y = x.to(parameters[\"device\"]), y.to(parameters[\"device\"])\n",
    "            optimizer.zero_grad()\n",
    "            x = x.view(x.size(0), -1)\n",
    "            outputs = model(x)\n",
    "            loss = criterion(outputs, y)\n",
    "\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            acc = (torch.sum(preds == y.data)) / len(x)\n",
    "\n",
    "            # Log trial metrics\n",
    "            run.log_metrics(\n",
    "                data={\n",
    "                    f\"trials/{trial}/metrics/batch/loss\": float(loss),\n",
    "                    f\"trials/{trial}/metrics/batch/acc\": float(acc),\n",
    "                },\n",
    "                step=step,\n",
    "            )\n",
    "\n",
    "            # Log best values across all trials\n",
    "            if best_acc is None or acc > best_acc:\n",
    "                best_acc = acc\n",
    "                run.log_configs(\n",
    "                    {\n",
    "                        \"best/trial\": trial,\n",
    "                        \"best/metrics/loss\": float(loss),\n",
    "                        \"best/metrics/acc\": float(acc),\n",
    "                        \"best/parameters/lr\": lr,\n",
    "                    }\n",
    "                )\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            step += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stop logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore the results in Neptune\n",
    "Follow the link to the run and explore the logged metadata in the Neptune app:\n",
    "\n",
    "- The best trial, with its metrics and parameters, is available in the *best* namespace\n",
    "- Metadata across all trials is available in the *trials* namespace\n",
    "\n",
    "To organize all relevant metadata in one view, create a [custom dashboard](https://docs-beta.neptune.ai/custom_dashboard). [See an example](https://scale.neptune.ai/o/examples/org/hpo/runs/details?viewId=9d4424ec-5c27-4933-9003-d62e0784ac68&detailsTab=dashboard&dashboardId=HPO-overview-9d4421e6-dfe4-400b-9dfb-d9b9e8a416b6&runIdentificationKey=HPO-11&type=run).\n",
    "\n",
    "To view best trials across different runs, you can also create [saved table views](https://docs-beta.neptune.ai/experiments_table#custom-views). [See an example](https://scale.neptune.ai/o/examples/org/hpo/runs/table?viewId=9d4424ec-5c27-4933-9003-d62e0784ac68&detailsTab=dashboard&dash=table&type=run)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log metadata from each HPO trial into separate runs\n",
    "\n",
    "You can also log metadata from each trial into separate runs. This way, you can track metadata from each trial separately.  \n",
    "Aggregated values can be logged to a parent sweep-level run. Sweep-level identifiers can be used to group all trials from the same sweep."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BaseModel(\n",
    "    input_size,\n",
    "    parameters[\"n_classes\"],\n",
    ").to(parameters[\"device\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a sweep-level identifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "sweep_id = str(uuid.uuid4())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize sweep-level run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_run = Run(run_id=f\"sweep-{sweep_id}\")\n",
    "\n",
    "sweep_run.add_tags([\"sweep\", \"notebook\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign sweep_id to sweep-level run as a group tag\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_run.add_tags([sweep_id], group_tags=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Log configuration common across all trials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in parameters:\n",
    "    if type(parameters[key]) not in ALLOWED_DATATYPES:\n",
    "        sweep_run.log_configs({f\"config/{key}\": str(parameters[key])})\n",
    "    else:\n",
    "        sweep_run.log_configs({f\"config/{key}\": parameters[key]})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize fields for best values across all trials\n",
    "best_acc = None\n",
    "\n",
    "for trial, lr in tqdm(\n",
    "    enumerate(learning_rates),\n",
    "    total=len(learning_rates),\n",
    "    desc=\"Trials\",\n",
    "):\n",
    "    # Create a trial-level run\n",
    "    with Run(run_id=f\"trial-{sweep_id}-{trial}\") as trial_run:\n",
    "        trial_run.add_tags([\"trial\", \"notebook\"])\n",
    "\n",
    "        # Add sweep_id to the trial-level run\n",
    "        trial_run.add_tags([sweep_id], group_tags=True)\n",
    "\n",
    "        # Log trial number and hyperparams\n",
    "        trial_run.log_configs({\"trial_num\": trial, \"parameters/lr\": lr})\n",
    "\n",
    "        optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "        step = 0\n",
    "\n",
    "        for epoch in trange(parameters[\"epochs\"], desc=f\"Trial {trial} - lr: {lr}\"):\n",
    "            trial_run.log_metrics(data={\"epochs\": epoch}, step=epoch)\n",
    "\n",
    "            for x, y in trainloader:\n",
    "                x, y = x.to(parameters[\"device\"]), y.to(parameters[\"device\"])\n",
    "                optimizer.zero_grad()\n",
    "                x = x.view(x.size(0), -1)\n",
    "                outputs = model(x)\n",
    "                loss = criterion(outputs, y)\n",
    "\n",
    "                _, preds = torch.max(outputs, 1)\n",
    "                acc = (torch.sum(preds == y.data)) / len(x)\n",
    "\n",
    "                # Log trial metrics\n",
    "                trial_run.log_metrics(\n",
    "                    data={\n",
    "                        \"metrics/batch/loss\": float(loss),\n",
    "                        \"metrics/batch/acc\": float(acc),\n",
    "                    },\n",
    "                    step=step,\n",
    "                )\n",
    "\n",
    "                # Log best values across all trials to Sweep-level run\n",
    "                if best_acc is None or acc > best_acc:\n",
    "                    best_acc = acc\n",
    "                    sweep_run.log_configs(\n",
    "                        {\n",
    "                            \"best/trial\": trial,\n",
    "                            \"best/metrics/loss\": float(loss),\n",
    "                            \"best/metrics/acc\": float(acc),\n",
    "                            \"best/parameters/lr\": lr,\n",
    "                        }\n",
    "                    )\n",
    "\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "\n",
    "                step += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stop the sweep-level run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_run.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore the results in Neptune\n",
    "Follow the link to the runs and explore the logged metadata in the Neptune app:\n",
    "\n",
    "- **Single run**\n",
    "  - The best trial, with its metrics and parameters, is available in the *best* namespace of the sweep-level run\n",
    "  - Metadata across all trials are available in the trial-level runs\n",
    "\n",
    "- **Multiple runs**\n",
    "  - To group all trials under a sweep, use the [run groups](https://docs-beta.neptune.ai/groups). [See an example](https://scale.neptune.ai/o/examples/org/hpo/runs/table?viewId=9d44261f-32a1-42e7-96ff-9b35edc4be66&detailsTab=dashboard&dash=table&type=run).\n",
    "  - To compare trails within or across sweeps, create a [multi-run dashboard](https://docs-beta.neptune.ai/custom_dashboard#multi-run-dashboard). [See an example](https://scale.neptune.ai/o/examples/org/hpo/runs/compare?viewId=9d44261f-32a1-42e7-96ff-9b35edc4be66&detailsTab=dashboard&dash=dashboard&dashboardId=Compare-trials-9d44284a-40fe-4614-a66d-a5ca81b8b4cd&type=run&compare=uIWrlI2f5Tyn_lrTzrCY6RSrOVUYtMkY0ozkGXHFv6E8). \n",
    "    - To compare the average of trials across different sweeps, turn on [*Average grouped runs*](https://docs-beta.neptune.ai/charts#comparing-grouped-runs) in the chart widget settings.\n",
    "  - To see both sweep-level and trial-level comparisons together, export charts or dashboards to a [report](https://docs-beta.neptune.ai/reports). [See an example](https://scale.neptune.ai/o/examples/org/hpo/reports/9d442900-19b4-47dc-a2e9-0faedc1f4d2c)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}