examples/xgraph/gradcam.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# GradCAM on ClinTox dataset for GIN"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from dig.xgraph.dataset import MoleculeDataset\n",
    "import torch\n",
    "from torch.utils.data import random_split\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.data import Data, InMemoryDataset, download_url, extract_zip\n",
    "from torch_geometric.data.dataset import files_exist\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",
   "source": [
    "### Load dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "def split_dataset(dataset, dataset_split=[0.8, 0.1, 0.1]):\n",
    "    dataset_len = len(dataset)\n",
    "    dataset_split = [int(dataset_len * dataset_split[0]),\n",
    "                     int(dataset_len * dataset_split[1]),\n",
    "                     0]\n",
    "    dataset_split[2] = dataset_len - dataset_split[0] - dataset_split[1]\n",
    "    train_set, val_set, test_set = \\\n",
    "        random_split(dataset, dataset_split)\n",
    "\n",
    "    return {'train': train_set, 'val': val_set, 'test': test_set}\n",
    "\n",
    "dataset = MoleculeDataset('datasets', 'clintox')\n",
    "dataset.data.x = dataset.data.x.to(torch.float32)\n",
    "dataset.data.y = dataset.data.y[:, 0] \n",
    "dim_node = dataset.num_node_features\n",
    "dim_edge = dataset.num_edge_features\n",
    "num_targets = dataset.num_classes\n",
    "num_classes = 2\n",
    "\n",
    "splitted_dataset = split_dataset(dataset)\n",
    "dataloader = DataLoader(splitted_dataset['test'], batch_size=1, shuffle=False)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load model and checkpoints"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dig.xgraph.models import GIN_3l\n",
    "\n",
    "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_3l(model_level='graph', dim_node=dim_node, dim_hidden=300, num_classes=num_classes)\n",
    "model.to(device)\n",
    "check_checkpoints()\n",
    "ckpt_path = osp.join('checkpoints', 'clintox', 'GIN_3l', '0', 'GIN_3l_best.ckpt')\n",
    "model.load_state_dict(torch.load(ckpt_path)['state_dict'])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Display example output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-3.7480,  3.3675]], device='cuda:0', grad_fn=<AddmmBackward>)\n"
     ]
    }
   ],
   "source": [
    "data = list(dataloader)[0].to(device)\n",
    "out = model(data.x, data.edge_index)\n",
    "print(out)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load the explainer"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "from dig.xgraph.method import GradCAM\n",
    "explainer = GradCAM(model, explain_graph=True)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Setup for evaluation"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Run explainer on the given model and dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "explain graph line 736\n",
      "explain graph line 924\n",
      "explain graph line 1197\n",
      "explain graph line 1276\n",
      "explain graph line 1322\n",
      "explain graph line 487\n",
      "explain graph line 1148\n",
      "explain graph line 1286\n",
      "explain graph line 438\n",
      "explain graph line 859\n",
      "explain graph line 395\n",
      "explain graph line 698\n",
      "explain graph line 658\n",
      "explain graph line 510\n",
      "explain graph line 842\n",
      "explain graph line 539\n",
      "explain graph line 1111\n",
      "explain graph line 469\n",
      "explain graph line 1220\n",
      "explain graph line 573\n",
      "explain graph line 679\n",
      "explain graph line 1337\n",
      "explain graph line 768\n",
      "explain graph line 613\n",
      "explain graph line 1424\n",
      "explain graph line 986\n",
      "explain graph line 233\n",
      "explain graph line 1182\n",
      "explain graph line 1317\n",
      "explain graph line 831\n",
      "explain graph line 290\n",
      "explain graph line 366\n",
      "explain graph line 309\n",
      "explain graph line 1165\n",
      "explain graph line 1254\n",
      "explain graph line 1207\n",
      "explain graph line 1294\n",
      "explain graph line 576\n",
      "explain graph line 861\n",
      "explain graph line 598\n",
      "explain graph line 406\n",
      "explain graph line 656\n",
      "explain graph line 1238\n",
      "explain graph line 780\n",
      "explain graph line 145\n",
      "explain graph line 1024\n",
      "explain graph line 778\n",
      "explain graph line 1088\n",
      "explain graph line 913\n",
      "explain graph line 152\n",
      "explain graph line 1290\n",
      "explain graph line 701\n",
      "explain graph line 871\n",
      "explain graph line 812\n",
      "explain graph line 528\n",
      "explain graph line 360\n",
      "explain graph line 1232\n",
      "explain graph line 824\n",
      "explain graph line 1463\n",
      "explain graph line 700\n",
      "explain graph line 179\n",
      "explain graph line 273\n",
      "explain graph line 1121\n",
      "explain graph line 398\n",
      "explain graph line 540\n",
      "explain graph line 228\n",
      "explain graph line 1102\n",
      "explain graph line 554\n",
      "explain graph line 119\n",
      "explain graph line 482\n",
      "explain graph line 312\n",
      "explain graph line 323\n",
      "explain graph line 1275\n",
      "explain graph line 682\n",
      "explain graph line 1241\n",
      "explain graph line 662\n",
      "explain graph line 299\n",
      "explain graph line 846\n",
      "explain graph line 954\n",
      "explain graph line 985\n",
      "explain graph line 126\n",
      "explain graph line 415\n",
      "explain graph line 78\n",
      "explain graph line 1430\n",
      "explain graph line 1119\n",
      "explain graph line 748\n",
      "explain graph line 560\n",
      "explain graph line 41\n",
      "explain graph line 1018\n",
      "explain graph line 738\n",
      "explain graph line 164\n",
      "explain graph line 988\n",
      "explain graph line 373\n",
      "explain graph line 1264\n",
      "explain graph line 798\n",
      "explain graph line 1412\n",
      "explain graph line 808\n",
      "explain graph line 229\n",
      "explain graph line 1300\n",
      "explain graph line 307\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Programs\\anaconda3\\envs\\main_env\\lib\\site-packages\\captum\\attr\\_utils\\gradient.py:31: UserWarning: Input Tensor 0 did not already require gradients, required_grads has been set automatically.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "for index, data in enumerate(dataloader):\n",
    "    print(f'explain graph line {dataloader.dataset.indices[index] + 2}')\n",
    "    data.to(device)\n",
    "\n",
    "    if torch.isnan(data.y[0].squeeze()):\n",
    "        continue\n",
    "\n",
    "    walks, masks, related_preds = \\\n",
    "        explainer(data.x, data.edge_index, sparsity=sparsity, num_classes=num_classes)\n",
    "\n",
    "    x_collector.collect_data(masks, related_preds, data.y[0].squeeze().long().item())\n",
    "\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",
    "\n",
    "    if index >= 99:\n",
    "        break"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Output metrics evaluation results"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fidelity: 0.4853\n",
      "Fidelity_inv: 0.0059\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}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}