examples/sslgraph/example_grace.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of GRACE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Node-level representation learning on CORA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dig.sslgraph.utils import Encoder\n",
    "from dig.sslgraph.evaluation import NodeUnsupervised\n",
    "from dig.sslgraph.dataset import get_node_dataset\n",
    "from dig.sslgraph.method import GRACE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loading dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = get_node_dataset('cora')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define encoders and contrastive model\n",
    "\n",
    "You can refer to [https://github.com/CRIPAC-DIG/GRACE/blob/master/config.yaml](https://github.com/CRIPAC-DIG/GRACE/blob/master/config.yaml) for detailed training configs.\n",
    "\n",
    "***Note***: Results in the GRACE paper uses different training-test splits to the public splits, due to which you may see different results in DIG and the original implementation of GRACE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 128\n",
    "encoder = Encoder(feat_dim=dataset[0].x.shape[1], hidden_dim=embed_dim, \n",
    "                  n_layers=2, gnn='gcn', node_level=True, graph_level=False)\n",
    "grace = GRACE(dim=embed_dim, dropE_rate_1=0.2, dropE_rate_2=0.4, \n",
    "              maskN_rate_1=0.3, maskN_rate_2=0.4, tau=0.4, device=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pretraining: epoch 2000: 100%|██████████| 2000/2000 [24:11<00:00,  1.38it/s, loss=7.392503]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best epoch 2000: acc 0.8212 (+/- 0.0044).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.8212000727653503"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluator = NodeUnsupervised(dataset, device=3, log_interval=100)\n",
    "evaluator.setup_train_config(p_lr=0.0005, p_epoch=2000, p_weight_decay=1e-5, comp_embed_on='cpu')\n",
    "evaluator.evaluate(learning_model=grace, 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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}