example_graphcl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of GraphCL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sslgraph.utils import Encoder, get_dataset\n",
    "from sslgraph.utils.eval_graph import EvalSemisupevised, EvalUnsupevised\n",
    "from sslgraph.contrastive.model import GraphCL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Semi-supervised learning on NCI1\n",
    "\n",
    "#### Load dataset\n",
    "\n",
    "In this example, we evaluate model on NCI1 dataset in the semi-supervised setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset, dataset_pretrain = get_dataset('NCI1', task='semisupervised')\n",
    "feat_dim = dataset[0].x.shape[1]\n",
    "embed_dim = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define your encoder and contrastive model (GraphCL)\n",
    "\n",
    "For semi-supervised setting, GraphCL uses ResGCN. \n",
    "\n",
    "Available augmentation includes: dropN, maskN, permE, subgraph, random[2-4]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = Encoder(feat_dim, embed_dim, n_layers=3, gnn='resgcn')\n",
    "graphcl = GraphCL(embed_dim, aug_1='subgraph', aug_2='subgraph')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define evaluatior instance\n",
    "\n",
    "In this example, we use a label rate of 1%.\n",
    "\n",
    "To setup configurations (num of epochs, learning rates, etc. for pretraining and finetuning), run\n",
    "\n",
    "\n",
    "`evaluator.setup_train_config(batch_size = 128,\n",
    "    p_optim = 'Adam', p_lr = 0.0001, p_weight_decay = 0, p_epoch = 100,\n",
    "    f_optim = 'Adam', f_lr = 0.001, f_weight_decay = 0, f_epoch = 100)`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator = EvalSemisupevised(dataset, dataset_pretrain, label_rate=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perform evaluation\n",
    "\n",
    "You can also perform evaluation with grid search on pre-training epoch and\n",
    "learning rate by running\n",
    "``\n",
    "evaluator.grid_search(learning_model=graphcl, encoder=encoder, \n",
    "    p_lr_lst=[0.1,0.01,0.001,0.0001], p_epoch_lst=[20,40,60,80,100])\n",
    "``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pretraining: epoch 100: 100%|██████████| 100/100 [14:51<00:00,  8.92s/it, loss=2.518394]\n",
      "Fold 1, finetuning: 100%|██████████| 100/100 [00:10<00:00,  9.64it/s, acc=0.6399, val_loss=3.3378]\n",
      "Fold 2, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.44it/s, acc=0.6229, val_loss=11.9820]\n",
      "Fold 3, finetuning: 100%|██████████| 100/100 [00:11<00:00,  9.00it/s, acc=0.5596, val_loss=2.5892]\n",
      "Fold 4, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.81it/s, acc=0.6010, val_loss=3.5205]\n",
      "Fold 5, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.66it/s, acc=0.6107, val_loss=4.9785]\n",
      "Fold 6, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.41it/s, acc=0.6691, val_loss=1.5976]\n",
      "Fold 7, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.81it/s, acc=0.5937, val_loss=2.1159]\n",
      "Fold 8, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.74it/s, acc=0.6058, val_loss=4.4338]\n",
      "Fold 9, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.86it/s, acc=0.6326, val_loss=2.1934]\n",
      "Fold 10, finetuning: 100%|██████████| 100/100 [00:11<00:00,  8.43it/s, acc=0.5888, val_loss=2.6012]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.6326034069061279, 0.046801406890153885)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluator.evaluate(learning_model=graphcl, encoder=encoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To reproduce results in the paper, you may want to perform grid search and run evaluation for 5 times and take the average."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Another example with a label rate of 10%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pretraining: epoch 100: 100%|██████████| 100/100 [10:36<00:00,  6.37s/it, loss=1.778023]\n",
      "Fold 1, finetuning: 100%|██████████| 100/100 [00:15<00:00,  6.26it/s, acc=0.7737, val_loss=0.7953]\n",
      "Fold 2, finetuning: 100%|██████████| 100/100 [00:15<00:00,  6.49it/s, acc=0.7397, val_loss=1.5799]\n",
      "Fold 3, finetuning: 100%|██████████| 100/100 [00:16<00:00,  6.04it/s, acc=0.7251, val_loss=1.4510]\n",
      "Fold 4, finetuning: 100%|██████████| 100/100 [00:15<00:00,  6.63it/s, acc=0.7567, val_loss=1.1257]\n",
      "Fold 5, finetuning: 100%|██████████| 100/100 [00:14<00:00,  6.82it/s, acc=0.7518, val_loss=1.2485]\n",
      "Fold 6, finetuning: 100%|██████████| 100/100 [00:16<00:00,  6.14it/s, acc=0.7324, val_loss=1.3859]\n",
      "Fold 7, finetuning: 100%|██████████| 100/100 [00:14<00:00,  6.79it/s, acc=0.7153, val_loss=1.6022]\n",
      "Fold 8, finetuning: 100%|██████████| 100/100 [00:15<00:00,  6.41it/s, acc=0.7251, val_loss=1.5585]\n",
      "Fold 9, finetuning: 100%|██████████| 100/100 [00:15<00:00,  6.43it/s, acc=0.7859, val_loss=1.1088]\n",
      "Fold 10, finetuning: 100%|██████████| 100/100 [00:16<00:00,  6.09it/s, acc=0.6959, val_loss=1.9115]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.7440389394760132, 0.031175630167126656)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = Encoder(feat_dim, embed_dim, n_layers=3, gnn='resgcn')\n",
    "graphcl = GraphCL(embed_dim, aug_1='random2', aug_2='random2')\n",
    "evaluator = EvalSemisupevised(dataset, dataset_pretrain, label_rate=0.1)\n",
    "evaluator.evaluate(learning_model=graphcl, encoder=encoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Unsupervised representation learning\n",
    "\n",
    "#### Load dataset\n",
    "\n",
    "In this example, we evaluate model on MUTAG dataset in the unsupervised setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = get_dataset('MUTAG', task='unsupervised')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define your encoder and contrastive model (GraphCL)\n",
    "\n",
    "For unsupervised setting, GraphCL uses GIN with jumping knowledge (with output_dim = hidden_dim * n_layers). \n",
    "\n",
    "Available augmentation includes: dropN, maskN, permE, subgraph, random[2-4]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 128\n",
    "encoder = Encoder(feat_dim=dataset[0].x.shape[1], hidden_dim=embed_dim, n_layers=3, gnn='gin')\n",
    "graphcl = GraphCL(embed_dim*3, aug_1=None, aug_2='random2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perform evaluation with grid search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pretraining: epoch 20: 100%|██████████| 20/20 [00:20<00:00,  1.04s/it, loss=5.231108]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best epoch 5: acc 0.8681 +/-(0.0751)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.8681286549707604, 0.07512400973969227)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluator = EvalUnsupevised(dataset, log_interval=5)\n",
    "evaluator.evaluate(learning_model=graphcl, encoder=encoder)"
   ]
  }
 ],
 "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": 4
}