examples/xgraph/gnnexplainer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# GNNExplainer on BA-Shapes dataset for 2-layer GIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from dig.xgraph.dataset import SynGraphDataset\n",
    "from dig.xgraph.models import *\n",
    "import torch\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.data import Data, InMemoryDataset, download_url, extract_zip\n",
    "import os.path as osp\n",
    "import os\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def index_to_mask(index, size):\n",
    "    mask = torch.zeros(size, dtype=torch.bool, device=index.device)\n",
    "    mask[index] = 1\n",
    "    return mask\n",
    "\n",
    "def split_dataset(dataset):\n",
    "    indices = []\n",
    "    num_classes = 4\n",
    "    train_percent = 0.7\n",
    "    for i in range(num_classes):\n",
    "        index = (dataset.data.y == i).nonzero().view(-1)\n",
    "        index = index[torch.randperm(index.size(0))]\n",
    "        indices.append(index)\n",
    "\n",
    "    train_index = torch.cat([i[:int(len(i) * train_percent)] for i in indices], dim=0)\n",
    "\n",
    "    rest_index = torch.cat([i[int(len(i) * train_percent):] for i in indices], dim=0)\n",
    "    rest_index = rest_index[torch.randperm(rest_index.size(0))]\n",
    "\n",
    "    dataset.data.train_mask = index_to_mask(train_index, size=dataset.data.num_nodes)\n",
    "    dataset.data.val_mask = index_to_mask(rest_index[:len(rest_index) // 2], size=dataset.data.num_nodes)\n",
    "    dataset.data.test_mask = index_to_mask(rest_index[len(rest_index) // 2:], size=dataset.data.num_nodes)\n",
    "\n",
    "    dataset.data, dataset.slices = dataset.collate([dataset.data])\n",
    "\n",
    "    return dataset\n",
    "\n",
    "dataset = SynGraphDataset('./datasets', 'BA_shapes')\n",
    "dataset.data.x = dataset.data.x.to(torch.float32)\n",
    "dataset.data.x = dataset.data.x[:, :1]\n",
    "# dataset.data.y = dataset.data.y[:, 2]\n",
    "dim_node = dataset.num_node_features\n",
    "dim_edge = dataset.num_edge_features\n",
    "# num_targets = dataset.num_classes\n",
    "num_classes = dataset.num_classes\n",
    "\n",
    "splitted_dataset = split_dataset(dataset)\n",
    "splitted_dataset.data.mask = splitted_dataset.data.test_mask\n",
    "splitted_dataset.slices['mask'] = splitted_dataset.slices['test_mask']\n",
    "dataloader = DataLoader(splitted_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load model and checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def check_checkpoints(root='./'):\n",
    "    if osp.exists(osp.join(root, 'checkpoints')):\n",
    "        return\n",
    "    url = ('https://github.com/divelab/DIG_storage/raw/main/xgraph/checkpoints.zip')\n",
    "    path = download_url(url, root)\n",
    "    extract_zip(path, root)\n",
    "    os.unlink(path)\n",
    "\n",
    "model = GIN_2l(model_level='node', dim_node=dim_node, dim_hidden=300, num_classes=num_classes)\n",
    "model.to(device)\n",
    "check_checkpoints()\n",
    "ckpt_path = osp.join('checkpoints', 'ba_shapes', 'GIN_2l', '0', 'GIN_2l_best.ckpt')\n",
    "model.load_state_dict(torch.load(ckpt_path)['state_dict'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Display example output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from dig.xgraph.method import GNNExplainer\n",
    "explainer = GNNExplainer(model, epochs=100, lr=0.01, explain_graph=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Setup for evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# --- Set the Sparsity to 0.5 ---\n",
    "sparsity = 0.5\n",
    "\n",
    "# --- Create data collector and explanation processor ---\n",
    "from dig.xgraph.evaluation import XCollector, ExplanationProcessor\n",
    "x_collector = XCollector(sparsity)\n",
    "# x_processor = ExplanationProcessor(model=model, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Run explainer on the given model and dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "explain graph 0 node 10\n",
      "explain graph 0 node 15\n",
      "explain graph 0 node 22\n",
      "explain graph 0 node 25\n",
      "explain graph 0 node 26\n",
      "explain graph 0 node 43\n",
      "explain graph 0 node 64\n",
      "explain graph 0 node 71\n",
      "explain graph 0 node 75\n",
      "explain graph 0 node 77\n",
      "explain graph 0 node 87\n",
      "explain graph 0 node 96\n",
      "explain graph 0 node 126\n",
      "explain graph 0 node 149\n",
      "explain graph 0 node 151\n",
      "explain graph 0 node 155\n",
      "explain graph 0 node 166\n",
      "explain graph 0 node 190\n",
      "explain graph 0 node 196\n",
      "explain graph 0 node 210\n",
      "explain graph 0 node 243\n",
      "explain graph 0 node 253\n",
      "explain graph 0 node 263\n",
      "explain graph 0 node 269\n",
      "explain graph 0 node 279\n",
      "explain graph 0 node 282\n",
      "explain graph 0 node 288\n",
      "explain graph 0 node 297\n",
      "explain graph 0 node 302\n",
      "explain graph 0 node 313\n",
      "explain graph 0 node 321\n",
      "explain graph 0 node 332\n",
      "explain graph 0 node 336\n",
      "explain graph 0 node 348\n",
      "explain graph 0 node 356\n",
      "explain graph 0 node 367\n",
      "explain graph 0 node 369\n",
      "explain graph 0 node 381\n",
      "explain graph 0 node 387\n",
      "explain graph 0 node 418\n",
      "explain graph 0 node 454\n",
      "explain graph 0 node 461\n",
      "explain graph 0 node 466\n",
      "explain graph 0 node 469\n",
      "explain graph 0 node 471\n",
      "explain graph 0 node 497\n",
      "explain graph 0 node 508\n",
      "explain graph 0 node 514\n",
      "explain graph 0 node 515\n",
      "explain graph 0 node 532\n",
      "explain graph 0 node 542\n",
      "explain graph 0 node 548\n",
      "explain graph 0 node 561\n",
      "explain graph 0 node 564\n",
      "explain graph 0 node 569\n",
      "explain graph 0 node 586\n",
      "explain graph 0 node 588\n",
      "explain graph 0 node 595\n",
      "explain graph 0 node 616\n",
      "explain graph 0 node 621\n",
      "explain graph 0 node 648\n",
      "explain graph 0 node 655\n",
      "explain graph 0 node 663\n",
      "explain graph 0 node 665\n",
      "explain graph 0 node 673\n",
      "explain graph 0 node 690\n",
      "explain graph 0 node 691\n",
      "explain graph 0 node 694\n",
      "explain graph 0 node 696\n",
      "explain graph 0 node 699\n"
     ]
    }
   ],
   "source": [
    "index = -1\n",
    "for i, data in enumerate(dataloader):\n",
    "   for j, node_idx in enumerate(torch.where(data.test_mask == True)[0]):\n",
    "        index += 1\n",
    "        print(f'explain graph {i} node {node_idx}')\n",
    "        data.to(device)\n",
    "\n",
    "        if torch.isnan(data.y[0].squeeze()):\n",
    "            continue\n",
    "    \n",
    "        edge_masks, hard_edge_masks, related_preds = \\\n",
    "            explainer(data.x, data.edge_index, sparsity=sparsity, num_classes=num_classes, node_idx=node_idx)\n",
    "    \n",
    "        x_collector.collect_data(hard_edge_masks, related_preds, data.y[0].squeeze().long().item())\n",
    "        # if you only have the edge masks without related_pred, please feed sparsity controlled mask to\n",
    "        # obtain the result: x_processor(data, masks, x_collector)\n",
    "        if index >= 99:\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Output metrics evaluation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fidelity: 0.3978\n",
      "Fidelity_inv: -0.0000\n",
      "Sparsity: 0.5000\n"
     ]
    }
   ],
   "source": [
    "print(f'Fidelity: {x_collector.fidelity:.4f}\\n'\n",
    "      f'Fidelity_inv: {x_collector.fidelity_inv:.4f}\\n'\n",
    "      f'Sparsity: {x_collector.sparsity:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}