hyperparameters search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "honest-civilization",
   "metadata": {},
   "source": [
    "### Hyperparameters Search for GCN Model\n",
    "\n",
    "- Prediction of aqueous solubility of small organic molecules using:\n",
    "    - Graph Convolutional Neural Networks (GCNs)\n",
    "    \n",
    "- Finding the best hyperparams for N=20 epochs, then train the best model.\n",
    "- Weight and Biases for logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "descending-diagnosis",
   "metadata": {},
   "outputs": [],
   "source": [
    "# libraries\n",
    "import os, math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix\n",
    "\n",
    "# pytorch\n",
    "import torch\n",
    "import torch.nn.functional as F \n",
    "from torch.nn import Linear\n",
    "from torch_geometric.nn import GCNConv, GATConv\n",
    "from torch_geometric.nn import global_mean_pool as gap\n",
    "from torch_geometric.nn import global_max_pool as gmp\n",
    "\n",
    "# from local files\n",
    "from custom_dataset import MoleculeDataset\n",
    "from model import GCN, GAT\n",
    "\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "active-terminology",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, n_features, hidden_channels, dropout_p=0.4):\n",
    "        super(GCN, self).__init__()\n",
    "        torch.manual_seed(21)\n",
    "        self.conv1 = GCNConv(n_features, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, int(hidden_channels))\n",
    "        self.conv3 = GCNConv(int(hidden_channels), int(hidden_channels))\n",
    "        self.linear = Linear(int(hidden_channels), 1)\n",
    "        self.dropout_p = dropout_p\n",
    "\n",
    "    def forward(self, data, edge_index, batch):\n",
    "        x, targets = data.x, data.y\n",
    "        # 1. Obtain the node embeddings\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "        \n",
    "        # 2. Aggregating message passing/embeddings\n",
    "        x = gap(x, batch)\n",
    "        \n",
    "        # 3. Apply the final classifier\n",
    "        x = F.dropout(x, p=self.dropout_p, training=self.training)\n",
    "\n",
    "        # model output from forward and loss \n",
    "        out = self.linear(x)  \n",
    "        loss = torch.nn.BCEWithLogitsLoss()(out, targets.reshape(-1, 1).type_as(out))\n",
    "\n",
    "        out = torch.sigmoid(out) # converting out proba in range [0, 1]\n",
    "        return out, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "automatic-awareness",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "batch_size=64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "noble-champagne",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GCN(\n",
      "  (conv1): GCNConv(30, 128)\n",
      "  (conv2): GCNConv(128, 128)\n",
      "  (conv3): GCNConv(128, 128)\n",
      "  (linear): Linear(in_features=128, out_features=1, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# model from model.py\n",
    "model = GCN(n_features=30, hidden_channels=128)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "excellent-monroe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device:  cpu\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Device: \", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "acute-atlanta",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "naughty-carter",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.login(key=\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "organizational-evans",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(project=\"\", entity=\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "understood-summer",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training and prediction with the model\n",
    "class Trainer:\n",
    "    \n",
    "    def __init__(self, model, optimizer, train_loader, valid_loader, batch_size=64):\n",
    "        self.model = model\n",
    "        self.train_loader = train_loader\n",
    "        self.valid_loader = valid_loader\n",
    "        self.batch_size = batch_size\n",
    "        self.optimizer = optimizer\n",
    "\n",
    "    # training model\n",
    "    def train_one_epoch(self, epoch):\n",
    "        self.model.train()\n",
    "\n",
    "        t_targets = []; p_targets = []; losses = []\n",
    "        tqdm_iter = tqdm(self.train_loader, total=len(self.train_loader))\n",
    "        for i, data in enumerate(tqdm_iter):\n",
    "\n",
    "            tqdm_iter.set_description(f\"Epoch {epoch}\")\n",
    "            self.optimizer.zero_grad()\n",
    "            outputs, loss = self.model(data, data.edge_index, data.batch)\n",
    "            targets = data.y\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "            y_true = self.process_output(targets)  # for one batch\n",
    "            y_proba = self.process_output(outputs.flatten()) # for one batch\n",
    "\n",
    "            auc = roc_auc_score(y_true, y_proba)\n",
    "            tqdm_iter.set_postfix(train_loss=loss.item(), train_auc=auc, valid_loss=None, valid_auc=None)\n",
    "            \n",
    "            losses.append(loss.item())\n",
    "            t_targets.extend(list(y_true))\n",
    "            p_targets.extend(list(y_proba))\n",
    "\n",
    "        epoch_auc = roc_auc_score(t_targets, p_targets)\n",
    "        epoch_loss = sum(losses)/len(losses)\n",
    "        return epoch_loss, epoch_auc, tqdm_iter\n",
    "\n",
    "\n",
    "    def process_output(self, out):\n",
    "        out = out.cpu().detach().numpy()\n",
    "        return out\n",
    "\n",
    "    \n",
    "    def validate_one_epoch(self, progress=None):\n",
    "        # model in eval model\n",
    "        self.model.eval()\n",
    "        \n",
    "        t_targets = []; p_targets = []; losses = []\n",
    "        for data in self.valid_loader:\n",
    "            \n",
    "            outputs, loss = self.model(data, data.edge_index, data.batch)\n",
    "            outputs, targets = outputs.flatten(), data.y\n",
    "            \n",
    "            y_proba = self.process_output(outputs)  # for one batch\n",
    "            y_true = self.process_output(targets) # for one batch \n",
    "            \n",
    "            t_targets.extend(list(y_true))\n",
    "            p_targets.extend(list(y_proba))\n",
    "            losses.append(loss.item())\n",
    "            \n",
    "        # progress stats\n",
    "        epoch_auc = roc_auc_score(t_targets, p_targets)\n",
    "        epoch_loss = sum(losses)/len(losses)\n",
    "        if progress:\n",
    "            progress_tracker = progress[\"tracker\"]\n",
    "            train_loss = progress[\"loss\"]\n",
    "            train_auc = progress[\"auc\"]\n",
    "            progress_tracker.set_postfix(train_loss=train_loss, train_auc=train_auc, valid_loss=epoch_loss, valid_auc=epoch_auc)              \n",
    "            progress_tracker.close()\n",
    "        return epoch_loss, epoch_auc\n",
    "            \n",
    "    # runs the training and validation trainer for n_epochs\n",
    "    def run(self, n_epochs=10):\n",
    "        \n",
    "        train_scores = []; train_losses = []\n",
    "        valid_scores = []; valid_losses = []\n",
    "        for e in range(1, n_epochs+1):\n",
    "            lt, at, progress_tracker = self.train_one_epoch(e)\n",
    "            \n",
    "            train_losses.append(lt)\n",
    "            train_scores.append(at)\n",
    "            \n",
    "            # validate this epoch\n",
    "            progress = {\"tracker\": progress_tracker, \"loss\": lt, \"auc\": at}  \n",
    "            lv, av = self.validate_one_epoch(progress=progress)  # pass training progress tracker to validation func\n",
    "            valid_losses.append(lv)\n",
    "            valid_scores.append(av)\n",
    "        \n",
    "        return (train_losses, train_scores), (valid_losses, valid_scores)\n",
    "            \n",
    "        \n",
    "    def predict(self, test_loader):\n",
    "        \n",
    "        self.model.eval()\n",
    "        predictions = []\n",
    "        tqdm_iter = tqdm(test_loader, total=len(test_loader))\n",
    "        for data in tqdm_iter:\n",
    "            with torch.no_grad():\n",
    "                o, _ = self.model(data, data.edge_index, data.batch)\n",
    "                o = self.process_output(o.flatten())\n",
    "                predictions.extend(list(o))\n",
    "            tqdm_iter.set_postfix(stage=\"test predictions\")\n",
    "        tqdm_iter.close()\n",
    "        return np.array(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "developed-polyester",
   "metadata": {},
   "source": [
    "#### Dataset: Custom molecular dataset (created in custom_dataset.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "entire-employer",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8847, 986)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset = MoleculeDataset(root=\"Dataset/\", filename=\"solubility-dataset-train.csv\")\n",
    "test_dataset = MoleculeDataset(root=\"Dataset/\", filename=\"solubility-dataset-test.csv\", test=True)\n",
    "\n",
    "# DataLoader\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "# data loaders size\n",
    "train_loader.dataset.length, test_loader.dataset.length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "nearby-encyclopedia",
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_config = {\n",
    "                  \"name\" : \"my-sweep\",\n",
    "                  \"method\" : \"bayes\",\n",
    "                  \"metric\": {\"name\": \"AUC\", \"goal\": \"maximize\"},\n",
    "                  \"parameters\" : {\n",
    "                    \"lr\": { \"min\": 0.01, \"max\": 0.1},\n",
    "                    \"weight_decay\": {\"min\": 1e-5, \"max\": 1e-2},\n",
    "                    \"hidden_channels\": {\"values\": [32, 64, 128, 256] },\n",
    "                    \"dropout_p\": {\"values\": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]},\n",
    "                    \"epochs\": {\"values\": [20]},  # training each models for only 20 epochs \n",
    "                  }\n",
    "            }\n",
    "\n",
    "sweep_id = wandb.sweep(sweep_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "domestic-trailer",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_true_test_labels(test_loader):\n",
    "    y_true_list = []\n",
    "    for data in test_loader:\n",
    "        ys = data.y.flatten()\n",
    "        y_true_list.extend(list(ys.cpu().detach().numpy()))\n",
    "    return np.array(y_true_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "medieval-watch",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tune_hyperparams(config=None):\n",
    "    N_EPOCH = 10\n",
    "    with wandb.init(config=config):\n",
    "        config = wandb.config\n",
    "        \n",
    "        # model prep and training\n",
    "        model = GCN(n_features=30, hidden_channels=config.hidden_channels, dropout_p=config.dropout_p)\n",
    "        model.to(device)\n",
    "        # Optimizer\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)\n",
    "        trainer = Trainer(model=model, optimizer=optimizer, train_loader=train_loader, valid_loader=test_loader)\n",
    "\n",
    "        (train_losses, train_scores), (valid_losses, valid_scores) = trainer.run(n_epochs=config.epochs)\n",
    "        \n",
    "        # logging hyperparameters\n",
    "        params = {}\n",
    "        for p, v in config.items():\n",
    "            params[p] = v\n",
    "        wandb.log(params)\n",
    "        \n",
    "        y_proba = trainer.predict(test_loader)\n",
    "        y_true = get_true_test_labels(test_loader=test_loader)\n",
    "        \n",
    "        # AUC on validation dataset\n",
    "        auc = roc_auc_score(y_true, y_proba)\n",
    "        wandb.log({\"AUC\": auc})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "advised-migration",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.agent(sweep_id, function=tune_hyperparams, count=50, project=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "maritime-enforcement",
   "metadata": {},
   "source": [
    "## Results from 50 Bayes search runs\n",
    "- Each run is trained only for 20 epochs.\n",
    "- Next, further train the best hyper-params combination for more epoches.\n",
    "\n",
    "\n",
    "![alt text](params_tuning.PNG \"Results from the 50 different searches\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "signed-corrections",
   "metadata": {},
   "source": [
    "- Bayes hyper-params tuner in WandB seems working! Only a few runs with very relatively small AUC at the beginning and most of the remaining runs towards the upper range of the AUC output.\n",
    "\n",
    "- Will train the best model (from the hyperparams combination above) for best possible performance in training notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dominican-foster",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}