W&B_Launch_Quickstart.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W&B Launch Quickstart",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/cvphelps/DeepLearningImplementations/blob/master/W%26B_Launch_Quickstart.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# W&B Launch Quickstart\n",
        "Use W&B Launch to quickly edit and re-run previous experiments."
      ],
      "metadata": {
        "id": "JUZQMIWsJx59"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3MEISScVOGA9"
      },
      "outputs": [],
      "source": [
        "!pip install wandb -qqq"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import wandb\n",
        "import math\n",
        "import random\n",
        "import torch, torchvision\n",
        "import torch.nn as nn\n",
        "import torchvision.transforms as T\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "def get_dataloader(is_train, batch_size, slice=5):\n",
        "    \"Get a training dataloader\"\n",
        "    full_dataset = torchvision.datasets.MNIST(root=\".\", train=is_train, transform=T.ToTensor(), download=True)\n",
        "    sub_dataset = torch.utils.data.Subset(full_dataset, indices=range(0, len(full_dataset), slice))\n",
        "    loader = torch.utils.data.DataLoader(dataset=sub_dataset, \n",
        "                                         batch_size=batch_size, \n",
        "                                         shuffle=True if is_train else False, \n",
        "                                         pin_memory=True, num_workers=2)\n",
        "    return loader\n",
        "\n",
        "def get_model(dropout):\n",
        "    \"A simple model\"\n",
        "    model = nn.Sequential(nn.Flatten(),\n",
        "                         nn.Linear(28*28, 256),\n",
        "                         nn.BatchNorm1d(256),\n",
        "                         nn.ReLU(),\n",
        "                         nn.Dropout(dropout),\n",
        "                         nn.Linear(256,10)).to(device)\n",
        "    return model\n",
        "\n",
        "def validate_model(model, valid_dl, loss_func, log_images=False, batch_idx=0):\n",
        "    \"Compute performance of the model on the validation dataset and log a wandb.Table\"\n",
        "    model.eval()\n",
        "    val_loss = 0.\n",
        "    with torch.inference_mode():\n",
        "        correct = 0\n",
        "        for i, (images, labels) in tqdm(enumerate(valid_dl), leave=False):\n",
        "            images, labels = images.to(device), labels.to(device)\n",
        "\n",
        "            # Forward pass ➡\n",
        "            outputs = model(images)\n",
        "            val_loss += loss_func(outputs, labels)*labels.size(0)\n",
        "\n",
        "            # Compute accuracy and accumulate\n",
        "            _, predicted = torch.max(outputs.data, 1)\n",
        "            correct += (predicted == labels).sum().item()\n",
        "\n",
        "            # Log one batch of images to the dashboard, always same batch_idx.\n",
        "            if i==batch_idx and log_images:\n",
        "                log_image_table(images, predicted, labels, outputs.softmax(dim=1))\n",
        "    return val_loss / len(valid_dl.dataset), correct / len(valid_dl.dataset)\n",
        "\n",
        "def log_image_table(images, predicted, labels, probs):\n",
        "    \"Log a wandb.Table with (img, pred, target, scores)\"\n",
        "    # 🐝 Create a wandb Table to log images, labels and predictions to\n",
        "    table = wandb.Table(columns=[\"image\", \"pred\", \"target\"]+[f\"score_{i}\" for i in range(10)])\n",
        "    for img, pred, targ, prob in zip(images.to(\"cpu\"), predicted.to(\"cpu\"), labels.to(\"cpu\"), probs.to(\"cpu\")):\n",
        "        table.add_data(wandb.Image(img[0].numpy()*255), pred, targ, *prob.numpy())\n",
        "    wandb.log({\"predictions_table\":table}, commit=False)"
      ],
      "metadata": {
        "id": "hEz6A8aSOaFC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 🐝 Initialise a wandb run\n",
        "wandb.init(\n",
        "    project=\"launch-quickstart\",\n",
        "    config={\n",
        "        \"epochs\": 10,\n",
        "        \"batch_size\": 128,\n",
        "        \"lr\": 1e-3,\n",
        "        \"dropout\": 0.5,\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Copy your config \n",
        "config = wandb.config\n",
        "\n",
        "# Get the data\n",
        "train_dl = get_dataloader(is_train=True, batch_size=config.batch_size)\n",
        "valid_dl = get_dataloader(is_train=False, batch_size=2*config.batch_size)\n",
        "n_steps_per_epoch = math.ceil(len(train_dl.dataset) / config.batch_size)\n",
        "\n",
        "# A simple MLP model\n",
        "model = get_model(config.dropout)\n",
        "\n",
        "# Make the loss and optimizer\n",
        "loss_func = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)\n",
        "\n",
        "# Training\n",
        "example_ct = 0\n",
        "step_ct = 0\n",
        "for epoch in tqdm(range(config.epochs)):\n",
        "    model.train()\n",
        "    for step, (images, labels) in enumerate(tqdm(train_dl, leave=False)):\n",
        "        images, labels = images.to(device), labels.to(device)\n",
        "\n",
        "        outputs = model(images)\n",
        "        train_loss = loss_func(outputs, labels)\n",
        "        optimizer.zero_grad()\n",
        "        train_loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        example_ct += len(images)\n",
        "        metrics = {\"train/train_loss\": train_loss, \n",
        "                    \"train/epoch\": (step + 1 + (n_steps_per_epoch * epoch)) / n_steps_per_epoch, \n",
        "                    \"train/example_ct\": example_ct}\n",
        "        \n",
        "        if step + 1 < n_steps_per_epoch:\n",
        "            # 🐝 Log train metrics to wandb \n",
        "            wandb.log(metrics)\n",
        "            \n",
        "        step_ct += 1\n",
        "\n",
        "    val_loss, accuracy = validate_model(model, valid_dl, loss_func, log_images=(epoch==(config.epochs-1)))\n",
        "\n",
        "    # 🐝 Log train and validation metrics to wandb\n",
        "    val_metrics = {\"val/val_loss\": val_loss, \n",
        "                    \"val/val_accuracy\": accuracy}\n",
        "    wandb.log({**metrics, **val_metrics})\n",
        "    \n",
        "    print(f\"Train Loss: {train_loss:.3f}, Valid Loss: {val_loss:3f}, Accuracy: {accuracy:.2f}\")\n",
        "\n",
        "    # If you had a test set, this is how you could log it as a Summary metric\n",
        "    wandb.summary['test_accuracy'] = 0.8\n",
        "\n",
        "# 🐝 Close your wandb run \n",
        "wandb.finish()"
      ],
      "metadata": {
        "id": "e6b0UpOmOH4m"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}