JNotebooks/W&B_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wandb tutorial\n",
    "\n",
    "What is wandb?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](wandb_screenshot.png \"Title\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchmetrics\n",
    "\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment tracking\n",
    "\n",
    "One of the most useful features of wandb is the ability to manage your experiments. We will first demonstrate basic usage of wandb to track our image classification results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set some variables\n",
    "batch_size = 4\n",
    "lr = 0.001\n",
    "epochs = 10\n",
    "optim_name = 'ADAM'\n",
    "loss_fn = 'Cross entropy'\n",
    "filters = 20\n",
    "\n",
    "# Generate our run id and config dictionary\n",
    "id = wandb.util.generate_id()\n",
    "\n",
    "config = {\n",
    "    \"batch_size\": batch_size, \n",
    "    \"learning_rate\": lr,\n",
    "    \"loss_function\": loss_fn,\n",
    "    \"optimizer\": optim_name,\n",
    "    \"run_id\": id,\n",
    "    \"epochs\": epochs,\n",
    "    \"filters\": filters,\n",
    "}\n",
    "\n",
    "# Set our run name, project, and anything else we want to initialize our run with\n",
    "run_name = f\"simple_CIFAR10_v0\"\n",
    "project=\"ENEL_645\"\n",
    "note = 'Simple CIFAR10 classification wandb tutorial'\n",
    "run = wandb.init(project=project, entity='natalia-dubljevic', id=id, name=run_name, config=config, notes=note)  # resume is True when resuming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download and deal with our data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.ToTensor()\n",
    "dataset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                       download=True, transform=transform)\n",
    "dataset = torch.utils.data.Subset(dataset, list(range(int(len(dataset) / 10))))\n",
    "\n",
    "trainset, valset = torch.utils.data.random_split(dataset, [0.7, 0.3])\n",
    "train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(valset, batch_size=1, shuffle=False)\n",
    "\n",
    "classes = ('airplane', 'automobile', 'bird', 'cat',\n",
    "           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_transform(img):\n",
    "    img = img.numpy()\n",
    "    img = np.transpose(img, (1, 2, 0))\n",
    "    return img\n",
    "\n",
    "# get some random training images\n",
    "dataiter = iter(train_loader)\n",
    "images, labels = next(dataiter)\n",
    "\n",
    "fig = plt.figure()\n",
    "for i, image in enumerate(images):\n",
    "    plt.subplot(1, 4, i+1)\n",
    "    plt.imshow(img_transform(image))\n",
    "    plt.title(classes[labels[i]])\n",
    "    plt.axis('off')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define our network architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, filters=20) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(3, filters, kernel_size=3, padding='same')\n",
    "        self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, padding='same')\n",
    "        self.conv3 = nn.Conv2d(filters, filters, kernel_size=3, padding='same')\n",
    "        \n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "        self.fc1 = nn.Linear(filters * 4 * 4, 120)\n",
    "        self.fc2 = nn.Linear(120, 32)\n",
    "        self.fc3 = nn.Linear(32, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(self.relu(self.conv1(x)))\n",
    "        x = self.pool(self.relu(self.conv2(x)))\n",
    "        x = self.pool(self.relu(self.conv3(x)))\n",
    "\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cuda check\n",
    "device='cuda:0' if torch.cuda.is_available() else 'cpu'\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Begin the training and validation loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net().to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=lr)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # Training loop\n",
    "    train_loss = 0.0\n",
    "    for i, data in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        imgs, labels = data[0].to(device), data[1].to(device)\n",
    "        outputs = net(imgs)\n",
    "\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    train_loss /= (i + 1)\n",
    "    print(f'train loss: {train_loss:.6f}', flush=True)\n",
    "\n",
    "    # Validation loop\n",
    "    val_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        imgs = []\n",
    "        preds = []\n",
    "        labels = []\n",
    "\n",
    "        # Create a table of predictions for later\n",
    "        columns = ['Image', 'Prediction', 'Truth']\n",
    "        for class_name in classes:\n",
    "            columns.append(\"score_\" + class_name)\n",
    "        predictions_table = wandb.Table(columns=columns)\n",
    "\n",
    "        for i, data in enumerate(val_loader):\n",
    "            img, label = data[0].to(device), data[1].to(device)\n",
    "            output = net(img)  # recall this is of size batch x 10\n",
    "            loss = criterion(output, label)\n",
    "            val_loss += loss.item()\n",
    "\n",
    "            _, pred = torch.max(output, 1)  # returns max value and index of max value\n",
    "            preds.append(pred.item())\n",
    "            labels.append(label.item())\n",
    "\n",
    "            if i in range(16):\n",
    "                # Create some lists to log some images\n",
    "                pred_class, label_class = classes[pred], classes[label]\n",
    "                imgs.append(wandb.Image(img, caption=f\"Pred: {pred_class}, Label: {label_class}\"))\n",
    "                \n",
    "                # Add data to our predictions table\n",
    "                row = [wandb.Image(img), pred_class, label_class]\n",
    "                output = nn.functional.softmax(torch.squeeze(output), dim=0)\n",
    "                for class_prob in output.tolist():\n",
    "                    row.append(np.round(class_prob, 4))\n",
    "                predictions_table.add_data(*row)\n",
    "\n",
    "    acc = torchmetrics.functional.accuracy(torch.Tensor(preds), torch.Tensor(labels), task=\"multiclass\", num_classes=10)\n",
    "\n",
    "    val_loss /= (i + 1)\n",
    "    print(f'val loss: {val_loss:.6f}', flush=True)\n",
    "    wandb.log({\"train_loss\": train_loss, \n",
    "                \"val_loss\": val_loss,\n",
    "                \"accuracy\": acc,\n",
    "                'img': imgs, \n",
    "                'prediction_table' : predictions_table,\n",
    "                \"conf_mat\" : wandb.plot.confusion_matrix(\n",
    "                    probs=None, y_true=labels, preds=preds,\n",
    "                    class_names=classes)},\n",
    "                step=epoch+1)\n",
    "    \n",
    "run.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameter sweeps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hyperparameter sweeps allow as to dial in on the most optimal hyperparameters for a given experiment. These tend to be more feasible for lighter models as they can be computationally expensive to run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](sweep_controller.png \"Title\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define our sweep configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_defaults = {\n",
    "    'loss_function': loss_fn,\n",
    "    'optimizer': optim_name,\n",
    "    'epochs': epochs,\n",
    "    'batch_size': batch_size\n",
    "}\n",
    "\n",
    "\n",
    "grid_sweep = {\n",
    "    'method': 'grid',  # or 'random' or 'bayes'\n",
    "    'name': 'grid_sweep',\n",
    "    'metric': {'goal': 'maximize', 'name': 'acc'},\n",
    "    'parameters': {\n",
    "        'filters': {'values': [10, 20, 30]},\n",
    "        'learning_rate': {'values': [0.01, 0.001]}\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "random_sweep = {\n",
    "    'method': 'random',\n",
    "    'name': 'random_sweep',\n",
    "    'metric': {'goal': 'maximize', 'name': 'acc'},\n",
    "    'parameters': {\n",
    "        'learning_rate': {\n",
    "           # val between exp(min) and exp(max) such that log is uniformly \n",
    "           #distributed between min and max\n",
    "            'distribution': 'log_uniform', \n",
    "            'min': -8, \n",
    "            'max': -5},\n",
    "        'filters': {\n",
    "            # a flat distribution between 0 and 0.1\n",
    "            'distribution': 'int_uniform',\n",
    "            'min': 10,\n",
    "            'max': 100}\n",
    "   }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up our training function for the sweep controller"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(config=None):\n",
    "    with wandb.init(config=config_defaults) as run:\n",
    "        # If called by wandb.agent, as below,\n",
    "        # this config will be set by Sweep Controller\n",
    "        config = wandb.config\n",
    "\n",
    "        net = Net(filters=config.filters).to(device)\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "        optimizer = optim.Adam(net.parameters(), lr=config.learning_rate)\n",
    "\n",
    "        for epoch in range(epochs):\n",
    "            train_loss = 0.0\n",
    "            for i, data in enumerate(train_loader):\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                imgs, labels = data[0].to(device), data[1].to(device)\n",
    "                outputs = net(imgs)\n",
    "                loss = criterion(outputs, labels)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "\n",
    "                train_loss += loss.item()\n",
    "\n",
    "            train_loss /= (i + 1)\n",
    "\n",
    "            val_loss = 0.0\n",
    "            with torch.no_grad():\n",
    "                imgs = []\n",
    "                preds = []\n",
    "                labels = []\n",
    "\n",
    "                for i, data in enumerate(val_loader):\n",
    "                    img, label = data[0].to(device), data[1].to(device)\n",
    "                    output = net(img)  # recall this is of size batch x 10\n",
    "                    loss = criterion(output, label)\n",
    "                    val_loss += loss.item()\n",
    "\n",
    "                    _, pred = torch.max(output, 1)  # returns max value and index of max value\n",
    "                    preds.append(pred)\n",
    "                    labels.append(label)\n",
    "\n",
    "                    if i in range(16):\n",
    "                        pred, label = classes[pred.item()], classes[label.item()]\n",
    "                        img = torch.squeeze(img.detach().cpu())\n",
    "                        imgs.append(wandb.Image(img, caption=f\"Pred: {pred}, Label: {label}\"))\n",
    "\n",
    "            acc = torchmetrics.functional.accuracy(torch.Tensor(preds), torch.Tensor(labels), task=\"multiclass\", num_classes=10)\n",
    "\n",
    "            val_loss /= (i + 1)\n",
    "            wandb.log({\"train_loss\": train_loss, \n",
    "                        \"val_loss\": val_loss,\n",
    "                        \"accuracy\": acc,\n",
    "                        'img': imgs}, step=epoch+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's try out a grid sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_id = wandb.sweep(grid_sweep, project=\"ENEL_645\", entity='natalia-dubljevic')\n",
    "wandb.agent(sweep_id, train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We can also do a random sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_id = wandb.sweep(random_sweep, project=\"ENEL_645\", entity='natalia-dubljevic')\n",
    "wandb.agent(sweep_id, train, count=6)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deep_env",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}