gdc_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import yaml\n",
    "import torch\n",
    "\n",
    "import scipy.sparse as sp\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from torch.optim import Adam, Optimizer\n",
    "from collections import defaultdict\n",
    "from torch_geometric.data import Data, InMemoryDataset\n",
    "\n",
    "from data import get_dataset, HeatDataset, PPRDataset, set_train_val_test_split\n",
    "from models import GCN\n",
    "from seeds import val_seeds, test_seeds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GCN with GDC\n",
    "This notebook demonstrates how to enhance GCN with GDC."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GDC\n",
    "At its core, preprocessing with GDC is just this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gdc(A: sp.csr_matrix, alpha: float, eps: float):\n",
    "    N = A.shape[0]\n",
    "\n",
    "    # Self-loops\n",
    "    A_loop = sp.eye(N) + A\n",
    "\n",
    "    # Symmetric transition matrix\n",
    "    D_loop_vec = A_loop.sum(0).A1\n",
    "    D_loop_vec_invsqrt = 1 / np.sqrt(D_loop_vec)\n",
    "    D_loop_invsqrt = sp.diags(D_loop_vec_invsqrt)\n",
    "    T_sym = D_loop_invsqrt @ A_loop @ D_loop_invsqrt\n",
    "\n",
    "    # PPR-based diffusion\n",
    "    S = alpha * sp.linalg.inv(sp.eye(N) - (1 - alpha) * T_sym)\n",
    "\n",
    "    # Sparsify using threshold epsilon\n",
    "    S_tilde = S.multiply(S >= eps)\n",
    "\n",
    "    # Column-normalized transition matrix on graph S_tilde\n",
    "    D_tilde_vec = S_tilde.sum(0).A1\n",
    "    T_S = S_tilde / D_tilde_vec\n",
    "    \n",
    "    return T_S"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Choose backend\n",
    "We will use the GPU in this notebook. If you want to use a CPU instead simply change this line to `cpu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load configuration\n",
    "The parameter settings for datasets and models as well as the training routine are stored in `config.yaml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('config.yaml', 'r') as c:\n",
    "    config = yaml.safe_load(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset and preprocess with GDC\n",
    "For convenience we will use a PyTorch Geometric InMemoryDataset in this notebook. `PPRDataset` (and `HeatDataset`) provide more flexibility and functionality than the above `gdc` method. However, their preprocessing is essentially the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {}\n",
    "\n",
    "for preprocessing in ['none', 'heat', 'ppr']:\n",
    "    if preprocessing == 'none':\n",
    "        dataset = get_dataset(\n",
    "            name=config['dataset_name'],\n",
    "            use_lcc=config['use_lcc']\n",
    "        )\n",
    "        dataset.data = dataset.data.to(device)\n",
    "        datasets[preprocessing] = dataset\n",
    "    elif preprocessing == 'heat':\n",
    "        dataset = HeatDataset(\n",
    "            name=config['dataset_name'],\n",
    "            use_lcc=config['use_lcc'],\n",
    "            t=config[preprocessing]['t'],\n",
    "            k=config[preprocessing]['k'],\n",
    "            eps=config[preprocessing]['eps']\n",
    "        )\n",
    "        dataset.data = dataset.data.to(device)\n",
    "        datasets[preprocessing] = dataset\n",
    "    elif preprocessing == 'ppr':\n",
    "        dataset = PPRDataset(\n",
    "            name=config['dataset_name'],\n",
    "            use_lcc=config['use_lcc'],\n",
    "            alpha=config[preprocessing]['alpha'],\n",
    "            k=config[preprocessing]['k'],\n",
    "            eps=config[preprocessing]['eps']\n",
    "        )\n",
    "        dataset.data = dataset.data.to(device)\n",
    "        datasets[preprocessing] = dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create GCN models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {}\n",
    "\n",
    "for preprocessing, dataset in datasets.items():\n",
    "    models[preprocessing] = GCN(\n",
    "        dataset,\n",
    "        hidden=config[preprocessing]['hidden_layers'] * [config[preprocessing]['hidden_units']],\n",
    "        dropout=config[preprocessing]['dropout']\n",
    "    ).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model: torch.nn.Module, optimizer: Optimizer, data: Data):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    logits = model(data)\n",
    "    loss = F.nll_loss(logits[data.train_mask], data.y[data.train_mask])\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model: torch.nn.Module, data: Data, test: bool):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        logits = model(data)\n",
    "    eval_dict = {}\n",
    "    keys = ['val', 'test'] if test else ['val']\n",
    "    for key in keys:\n",
    "        mask = data[f'{key}_mask']\n",
    "        # loss = F.nll_loss(logits[mask], data.y[mask]).item()\n",
    "        # eval_dict[f'{key}_loss'] = loss\n",
    "        pred = logits[mask].max(1)[1]\n",
    "        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "        eval_dict[f'{key}_acc'] = acc\n",
    "    return eval_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(dataset: InMemoryDataset,\n",
    "        model: torch.nn.Module,\n",
    "        seeds: np.ndarray,\n",
    "        test: bool = False,\n",
    "        max_epochs: int = 10000,\n",
    "        patience: int = 100,\n",
    "        lr: float = 0.01,\n",
    "        weight_decay: float = 0.01,\n",
    "        num_development: int = 1500,\n",
    "        device: str = 'cuda'):\n",
    "    start_time = time.perf_counter()\n",
    "\n",
    "    best_dict = defaultdict(list)\n",
    "\n",
    "    cnt = 0\n",
    "    for seed in tqdm(seeds):\n",
    "        dataset.data = set_train_val_test_split(\n",
    "            seed,\n",
    "            dataset.data,\n",
    "            num_development=num_development,\n",
    "        ).to(device)\n",
    "        model.to(device).reset_parameters()\n",
    "        optimizer = Adam(\n",
    "            [\n",
    "                {'params': model.non_reg_params, 'weight_decay': 0},\n",
    "                {'params': model.reg_params, 'weight_decay': weight_decay}\n",
    "            ],\n",
    "            lr=lr\n",
    "        )\n",
    "\n",
    "        patience_counter = 0\n",
    "        tmp_dict = {'val_acc': 0}\n",
    "\n",
    "        for epoch in range(1, max_epochs + 1):\n",
    "            if patience_counter == patience:\n",
    "                break\n",
    "\n",
    "            train(model, optimizer, dataset.data)\n",
    "            eval_dict = evaluate(model, dataset.data, test)\n",
    "\n",
    "            if eval_dict['val_acc'] < tmp_dict['val_acc']:\n",
    "                patience_counter += 1\n",
    "            else:\n",
    "                patience_counter = 0\n",
    "                tmp_dict['epoch'] = epoch\n",
    "                for k, v in eval_dict.items():\n",
    "                    tmp_dict[k] = v\n",
    "\n",
    "        for k, v in tmp_dict.items():\n",
    "            best_dict[k].append(v)\n",
    "            \n",
    "    best_dict['duration'] = time.perf_counter() - start_time\n",
    "    return dict(best_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are training the models 100 times on different splits, so this will take a couple of minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d40a57c40dbd4d5abbeaa168e9cecca3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c982d5582ab94a6f9a74de84dfe8a319",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7aee5d03a8154762bc17b807ec7d3e40",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results = {}\n",
    "\n",
    "for preprocessing in ['none', 'heat', 'ppr']:\n",
    "    results[preprocessing] = run(\n",
    "        datasets[preprocessing],\n",
    "        models[preprocessing],\n",
    "        seeds=test_seeds if config['test'] else val_seeds,\n",
    "        lr=config[preprocessing]['lr'],\n",
    "        weight_decay=config[preprocessing]['weight_decay'],\n",
    "        test=config['test'],\n",
    "        num_development=config['num_development'],\n",
    "        device=device\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate statistics using bootstrapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _, best_dict in results.items():\n",
    "    boots_series = sns.algorithms.bootstrap(best_dict['val_acc'], func=np.mean, n_boot=1000)\n",
    "    best_dict['val_acc_ci'] = np.max(np.abs(sns.utils.ci(boots_series, 95) - np.mean(best_dict['val_acc'])))\n",
    "    if 'test_acc' in best_dict:\n",
    "        boots_series = sns.algorithms.bootstrap(best_dict['test_acc'], func=np.mean, n_boot=1000)\n",
    "        best_dict['test_acc_ci'] = np.max(\n",
    "            np.abs(sns.utils.ci(boots_series, 95) - np.mean(best_dict['test_acc']))\n",
    "        )\n",
    "\n",
    "    for k, v in best_dict.items():\n",
    "        if 'acc_ci' not in k and k != 'duration':\n",
    "            best_dict[k] = np.mean(best_dict[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "none: Mean accuracy: 81.74 +- 0.25%\n",
      "heat: Mean accuracy: 83.47 +- 0.21%\n",
      "ppr: Mean accuracy: 83.64 +- 0.23%\n"
     ]
    }
   ],
   "source": [
    "for preprocessing in ['none', 'heat', 'ppr']:\n",
    "    mean_acc = results[preprocessing]['test_acc']\n",
    "    uncertainty = results[preprocessing]['test_acc_ci']\n",
    "    print(f\"{preprocessing}: Mean accuracy: {100 * mean_acc:.2f} +- {100 * uncertainty:.2f}%\")"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}