Launch_Conditional_Config_Sweeps_(+_multi_objective_optimization).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "cf8fe4925a984e4583b8a39e8ba7943f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "VBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "VBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7972b301b06440d7a6e54698f596576f",
              "IPY_MODEL_b4bb913cc0e44db1bfc3ac241fc850ac"
            ],
            "layout": "IPY_MODEL_7970871ca8614c13978cd26fb30c762f"
          }
        },
        "7972b301b06440d7a6e54698f596576f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "LabelModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "LabelModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "LabelView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8913122c01d345f790029d640c4fc147",
            "placeholder": "​",
            "style": "IPY_MODEL_4cb77445e94845e1814f9806d1dee76b",
            "value": "Waiting for wandb.init()...\r"
          }
        },
        "b4bb913cc0e44db1bfc3ac241fc850ac": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f1e9a1ba48d94251ad991919a53d626e",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_30d37dcffef14d068b383244040107e2",
            "value": 1
          }
        },
        "7970871ca8614c13978cd26fb30c762f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8913122c01d345f790029d640c4fc147": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4cb77445e94845e1814f9806d1dee76b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f1e9a1ba48d94251ad991919a53d626e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "30d37dcffef14d068b383244040107e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jonbaer/googlecolab/blob/master/Launch_Conditional_Config_Sweeps_(%2B_multi_objective_optimization).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conditional Configuration of Hyperparameter Search\n",
        "\n",
        "W&B Launch allows you to express more sophisticated hyperparameter search logic, using libraries like [Optuna](https://optuna.org/#key_features) along with Python definitions for search logic.   This allows you to find optimal hyperparameters more efficiently--and it's particularly powerful when paired with W&B's ability to push searches to powerful compute clusters.\n",
        "\n",
        "In this tutorial, you'll:\n",
        "\n",
        "\n",
        "* Define conditional search logic with one level of nesting\n",
        "* Save that logic to W&B\n",
        "* Define a sweep configuration using Optuna\n",
        "* Launch a sweep applying that logic to find optimal hyperparameters.\n",
        "\n",
        "From here, you'll be positioned to massively scale up the sophistication and scale of your hyperparameter sweeps to improve your models.\n",
        "\n",
        "Before you start, **make sure you have a Launch queue up and an agent running** (or follow [this guide](https://colab.research.google.com/drive/1wX0OSVxZJDHRsZaOaOEDx-lLUrO1hHgP#scrollTo=jhm3qUUxk69o&forceEdit=true&sandboxMode=true) to create one).  This is what will execute the sweep you define.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "AFEzIxA6foC7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Get set up with W&B"
      ],
      "metadata": {
        "id": "jtK9hSRIioAr"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UWoFnvBGKIgQ"
      },
      "outputs": [],
      "source": [
        "# Install W&B\n",
        "!pip install wandb>=0.15.8 -qqq"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Log in to your free W&B Account\n",
        "import wandb\n",
        "import inspect\n",
        "import yaml\n",
        "\n",
        "wandb.login()\n"
      ],
      "metadata": {
        "id": "SEjC8XAwHgeT",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "1db206f7-7d77-4af2-b711-8ece1af595b5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "        window._wandbApiKey = new Promise((resolve, reject) => {\n",
              "            function loadScript(url) {\n",
              "            return new Promise(function(resolve, reject) {\n",
              "                let newScript = document.createElement(\"script\");\n",
              "                newScript.onerror = reject;\n",
              "                newScript.onload = resolve;\n",
              "                document.body.appendChild(newScript);\n",
              "                newScript.src = url;\n",
              "            });\n",
              "            }\n",
              "            loadScript(\"https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js\").then(() => {\n",
              "            const iframe = document.createElement('iframe')\n",
              "            iframe.style.cssText = \"width:0;height:0;border:none\"\n",
              "            document.body.appendChild(iframe)\n",
              "            const handshake = new Postmate({\n",
              "                container: iframe,\n",
              "                url: 'https://wandb.ai/authorize'\n",
              "            });\n",
              "            const timeout = setTimeout(() => reject(\"Couldn't auto authenticate\"), 5000)\n",
              "            handshake.then(function(child) {\n",
              "                child.on('authorize', data => {\n",
              "                    clearTimeout(timeout)\n",
              "                    resolve(data)\n",
              "                });\n",
              "            });\n",
              "            })\n",
              "        });\n",
              "    "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "PROJECT = \"fmnist_job_tutorial\"\n",
        "ENTITY = \"\" ## Put in your entity\n",
        "QUEUE = \"tutorial-run-queue\" ## Put in a Launch queue you've created and started"
      ],
      "metadata": {
        "id": "lvv0TrFjwVZ3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2) Define conditional hyperparameter search logic\n",
        "\n",
        "Here we're doing just one level of nesting, making the batch size conditional on the `train_size` setting.  But you can make the search space arbitrarily complex using Python."
      ],
      "metadata": {
        "id": "NMVZsthv2hcS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def objective(trial):\n",
        "    \"\"\"Optuna objective function to create a config using pythonic search spaces.\n",
        "​\n",
        "    Does not actually train, but is logged to wandb an artifact and used in the\n",
        "    Optuna sweeps on launch scheduler.\"\"\"\n",
        "\n",
        "    train_size = trial.suggest_categorical('train_size', ['small', 'medium', 'large'])\n",
        "\n",
        "    if train_size == 'small':\n",
        "        batch_size = trial.suggest_int('batch_size', 16, 64)\n",
        "    elif train_size == 'medium':\n",
        "        batch_size = trial.suggest_int('batch_size', 8, 32)\n",
        "    else:\n",
        "        batch_size = trial.suggest_int('batch_size', 2, 16)\n",
        "\n",
        "    lr = trial.suggest_float('learning_rate', 0.000001, 1.0)\n",
        "\n",
        "    sleep = trial.suggest_float('sleep', 0.3, 1.0)\n",
        "\n",
        "    a = trial.suggest_float('a', 0, 10)\n",
        "\n",
        "    print(f\"{train_size=} {batch_size=} {lr=} {sleep=} {a=}\")\n",
        "\n",
        "    # !! don't actually train, return -1\n",
        "    return -1"
      ],
      "metadata": {
        "id": "cjovQwWjwQah"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3) Test the conditional configuration function.\n",
        "\n",
        "We'll import Optuna to test the function we just constructed before we save it to W&B.\n",
        "\n",
        "  "
      ],
      "metadata": {
        "id": "ItnN7VP23Cpb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install optuna"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tprlr3r6x1y9",
        "outputId": "f7dfa78d-ec3c-4349-dc27-758cec46343e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting optuna\n",
            "  Downloading optuna-3.3.0-py3-none-any.whl (404 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m404.2/404.2 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting alembic>=1.5.0 (from optuna)\n",
            "  Downloading alembic-1.12.0-py3-none-any.whl (226 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m226.0/226.0 kB\u001b[0m \u001b[31m13.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting cmaes>=0.10.0 (from optuna)\n",
            "  Downloading cmaes-0.10.0-py3-none-any.whl (29 kB)\n",
            "Collecting colorlog (from optuna)\n",
            "  Downloading colorlog-6.7.0-py2.py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from optuna) (1.23.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from optuna) (23.2)\n",
            "Requirement already satisfied: sqlalchemy>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from optuna) (2.0.21)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from optuna) (4.66.1)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from optuna) (6.0.1)\n",
            "Collecting Mako (from alembic>=1.5.0->optuna)\n",
            "  Downloading Mako-1.2.4-py3-none-any.whl (78 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.7/78.7 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: typing-extensions>=4 in /usr/local/lib/python3.10/dist-packages (from alembic>=1.5.0->optuna) (4.5.0)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy>=1.3.0->optuna) (3.0.0)\n",
            "Requirement already satisfied: MarkupSafe>=0.9.2 in /usr/local/lib/python3.10/dist-packages (from Mako->alembic>=1.5.0->optuna) (2.1.3)\n",
            "Installing collected packages: Mako, colorlog, cmaes, alembic, optuna\n",
            "Successfully installed Mako-1.2.4 alembic-1.12.0 cmaes-0.10.0 colorlog-6.7.0 optuna-3.3.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "    import optuna\n",
        "\n",
        "    study = optuna.create_study(direction='maximize')\n",
        "    study.optimize(objective, n_trials=2)\n",
        "    print(f\"Best trial: {study.best_trial.value=} {study.best_params=}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K0iY3_NhxlVE",
        "outputId": "3066f28f-59df-4c68-d969-52b0e683fddc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[I 2023-10-11 17:28:32,961] A new study created in memory with name: no-name-634f909f-5dd4-4fe1-93d5-43c01307d0f7\n",
            "[I 2023-10-11 17:28:32,990] Trial 0 finished with value: -1.0 and parameters: {'train_size': 'small', 'batch_size': 21, 'learning_rate': 0.5094918743625853, 'sleep': 0.9445360913643575, 'a': 1.059722445186102}. Best is trial 0 with value: -1.0.\n",
            "[I 2023-10-11 17:28:33,002] Trial 1 finished with value: -1.0 and parameters: {'train_size': 'large', 'batch_size': 2, 'learning_rate': 0.9278913345116473, 'sleep': 0.9043828540124754, 'a': 0.9869390001700007}. Best is trial 0 with value: -1.0.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "train_size='small' batch_size=21 lr=0.5094918743625853 sleep=0.9445360913643575 a=1.059722445186102\n",
            "train_size='large' batch_size=2 lr=0.9278913345116473 sleep=0.9043828540124754 a=0.9869390001700007\n",
            "Best trial: study.best_trial.value=-1.0 study.best_params={'train_size': 'small', 'batch_size': 21, 'learning_rate': 0.5094918743625853, 'sleep': 0.9445360913643575, 'a': 1.059722445186102}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 4) Save the configuration to W&B as an artifact.\n",
        "\n",
        "Now let's save the conditional search logic to W&B as an artifact."
      ],
      "metadata": {
        "id": "cvZ8oeQ23RnD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ARTIFACT_FILENAME = \"optuna_wandb.py\"\n",
        "ARTIFACT_NAME = \"optuna-config\"\n",
        "\n",
        "\"\"\"write function to its own file\"\"\"\n",
        "function_lines = inspect.getsource(objective)\n",
        "with open(ARTIFACT_FILENAME, 'w') as f:\n",
        "    f.write(function_lines)\n",
        "\n",
        "\"\"\"create and log artifact to wandb\"\"\"\n",
        "run = wandb.init(project=PROJECT, entity=ENTITY)\n",
        "artifact = run.log_artifact(ARTIFACT_FILENAME, name=ARTIFACT_NAME, type='optuna')\n",
        "run.finish()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 273,
          "referenced_widgets": [
            "cf8fe4925a984e4583b8a39e8ba7943f",
            "7972b301b06440d7a6e54698f596576f",
            "b4bb913cc0e44db1bfc3ac241fc850ac",
            "7970871ca8614c13978cd26fb30c762f",
            "8913122c01d345f790029d640c4fc147",
            "4cb77445e94845e1814f9806d1dee76b",
            "f1e9a1ba48d94251ad991919a53d626e",
            "30d37dcffef14d068b383244040107e2"
          ]
        },
        "id": "6ApCZIh2x8ZV",
        "outputId": "c103f355-44db-4013-df33-a77f1729f839"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mtristan-spaulding\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.01111422133333488, max=1.0)…"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "cf8fe4925a984e4583b8a39e8ba7943f"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Tracking run with wandb version 0.15.12"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Run data is saved locally in <code>/content/wandb/run-20231011_172836-ar8rf9d9</code>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Syncing run <strong><a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/runs/ar8rf9d9' target=\"_blank\">classic-violet-5</a></strong> to <a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              " View project at <a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial' target=\"_blank\">https://wandb.ai/tristan-spaulding/fmnist_job_tutorial</a>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              " View run at <a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/runs/ar8rf9d9' target=\"_blank\">https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/runs/ar8rf9d9</a>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              " View run <strong style=\"color:#cdcd00\">classic-violet-5</strong> at: <a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/runs/ar8rf9d9' target=\"_blank\">https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/runs/ar8rf9d9</a><br/> View job at <a href='https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/jobs/QXJ0aWZhY3RDb2xsZWN0aW9uOjEwNTk5Nzc3OA==/version_details/v0' target=\"_blank\">https://wandb.ai/tristan-spaulding/fmnist_job_tutorial/jobs/QXJ0aWZhY3RDb2xsZWN0aW9uOjEwNTk5Nzc3OA==/version_details/v0</a><br/>Synced 4 W&B file(s), 0 media file(s), 5 artifact file(s) and 1 other file(s)"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Find logs at: <code>./wandb/run-20231011_172836-ar8rf9d9/logs</code>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 5) Define a sweep configuration\n",
        "\n",
        "Now that we've logged the condition configuration artifact, we'll define a configuration for our sweep.\n",
        "\n",
        "(To adapt this for a different job, substitute your job for the top-level `job` field; here we'll use an [FMNIST job](https://colab.research.google.com/drive/1wX0OSVxZJDHRsZaOaOEDx-lLUrO1hHgP#scrollTo=tVR5Zn90h0pP).)\n",
        "\n",
        "We'll use a custom Optuna scheduler."
      ],
      "metadata": {
        "id": "TV8liA2a15g7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config = {\n",
        "    \"metric\": {\"name\": \"epoch/val_loss\", \"goal\": \"minimize\"},\n",
        "    \"run_cap\": 4,\n",
        "    \"job\": \"wandb/jobs/FMNIST Training:latest\",\n",
        "    \"scheduler\": {\n",
        "        \"job\": \"wandb/sweep-jobs/job-optuna-sweep-scheduler:latest\",\n",
        "        \"num_workers\": 2,\n",
        "        \"settings\": {\n",
        "            \"optuna_source\": f\"{ENTITY}/{PROJECT}/{artifact.wait().name}\",\n",
        "            \"optuna_source_filename\": ARTIFACT_FILENAME,\n",
        "            # optional sampler args\n",
        "            \"pruner\": {\n",
        "                \"type\": \"PercentilePruner\",\n",
        "                \"args\": {\n",
        "                    \"percentile\": 0.25,\n",
        "                    \"n_startup_trials\": 2,\n",
        "                    \"n_min_trials\": 1,  # min epochs before prune\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        "    # parameters are not needed when loading a conditional config from an artifact\n",
        "    # \"parameters\": {\n",
        "    #     'epochs': {'values': [5, 10, 15]},\n",
        "    #     'lr': {'max': 0.1, 'min': 0.0001}\n",
        "    # }\n",
        "}\n",
        "\n",
        "# write config to file\n",
        "config_filename = \"sweep-config.yaml\"\n",
        "yaml.dump(config, open(config_filename, \"w\"))"
      ],
      "metadata": {
        "id": "UrYv9vL_yOGi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#6) Launch the sweep!\n",
        "\n",
        "We can now launch the sweep and see its progress in W&B.\n"
      ],
      "metadata": {
        "id": "FytXg_eO3gl8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! wandb launch-sweep sweep-config.yaml -e $ENTITY -p $PROJECT -q $QUEUE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sHcr30CKybN7",
        "outputId": "cf86e5b9-4744-43fe-cf2d-7a075f35e933"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Using a scheduler job for launch sweeps is *experimental* and may change without warning\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m:   2 of 2 files downloaded.  \n",
            "\u001b[34m\u001b[1mwandb\u001b[0m:   2 of 2 files downloaded.  \n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[35mlaunch:\u001b[0m Launching run into tristan-spaulding/fmnist_job_tutorial\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Malformed sweep config detected! This may cause your sweep to behave in unexpected ways.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m To avoid this, please fix the sweep config schema violations below:\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m   Violation 1. 'parameters' is a required property\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Created sweep with ID: \u001b[33mtx33s6tt\u001b[0m\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: View sweep at: \u001b[34m\u001b[4mhttps://wandb.ai/tristan-spaulding/fmnist_job_tutorial/sweeps/tx33s6tt\u001b[0m\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Scheduler added to launch queue (tutorial-run-queue2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "🖥 Go take a look at your terminal running the launch agent to see the sweep runs running."
      ],
      "metadata": {
        "id": "ErJ0basOCxlA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##7) (Bonus) Create a multi-objective sweep\n",
        "\n",
        "Optuna supports optimizing over multiple objectives at once, so in this example we will define a hyperparameter sweep to minimize validation loss, while simultaneously maximizing the categorical accuracy."
      ],
      "metadata": {
        "id": "ENxAqpK7E_X3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config = {\n",
        "    \"run_cap\": 4,\n",
        "    \"job\": \"wandb/jobs/FMNIST Training:latest\",\n",
        "    \"scheduler\": {\n",
        "        \"job\": \"wandb/sweep-jobs/job-optuna-sweep-scheduler:latest\",\n",
        "        \"num_workers\": 2,\n",
        "        \"settings\": {\n",
        "            \"metrics\": [\n",
        "              {\"name\": \"epoch/val_loss\", \"goal\": \"minimize\"},\n",
        "              {\"name\": \"epoch/val_sparse_categorical_accuracy\", \"goal\": \"maximize\"}\n",
        "            ]\n",
        "        }\n",
        "    },\n",
        "    \"parameters\": {\n",
        "        'epochs': {'values': [2, 5, 10, 15]},\n",
        "        'lr': {'max': 0.1, 'min': 0.0001}\n",
        "    }\n",
        "}\n",
        "\n",
        "# write config to file\n",
        "config_filename = \"sweep-config-multi.yaml\"\n",
        "yaml.dump(config, open(config_filename, \"w\"))\n",
        "\n",
        "# and launch the sweep\n",
        "! wandb launch-sweep sweep-config-multi.yaml -e $ENTITY -p $PROJECT -q $QUEUE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c8AVhsYKE-PY",
        "outputId": "1ef928ca-c5ff-4148-e03e-91f33c011bee"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Using a scheduler job for launch sweeps is *experimental* and may change without warning\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m:   2 of 2 files downloaded.  \n",
            "\u001b[34m\u001b[1mwandb\u001b[0m:   2 of 2 files downloaded.  \n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[35mlaunch:\u001b[0m Launching run into tristan-spaulding/fmnist_job_tutorial\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Created sweep with ID: \u001b[33mk4jukhtr\u001b[0m\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: View sweep at: \u001b[34m\u001b[4mhttps://wandb.ai/tristan-spaulding/fmnist_job_tutorial/sweeps/k4jukhtr\u001b[0m\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Scheduler added to launch queue (tutorial-run-queue2)\n"
          ]
        }
      ]
    }
  ]
}