examples/xgraph/deeplift.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# DeepLIFT on Tox21 dataset for GCN"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/shurui.gui/anaconda3/envs/torch_110/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import random_split\n",
    "from torch_geometric.data import download_url, extract_zip\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "from dig.xgraph.dataset import MoleculeDataset\n",
    "from dig.xgraph.models import GCN_3l\n",
    "from dig.xgraph.utils.compatibility import compatible_state_dict\n",
    "from dig.xgraph.utils.init import fix_random_seed\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 = random_split(dataset, dataset_split)\n",
    "\n",
    "    return {'train': train_set, 'val': val_set, 'test': test_set}\n",
    "\n",
    "\n",
    "fix_random_seed(123)\n",
    "dataset = MoleculeDataset('datasets', 'Tox21')\n",
    "dataset.data.x = dataset.data.x.to(torch.float32)\n",
    "dataset.data.y = dataset.data.y[:, 2]  # the target 2 task.\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": [
    "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",
    "\n",
    "model = GCN_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', 'tox21', 'GCN_3l', '2', 'GCN_3l_best.ckpt')\n",
    "state_dict = compatible_state_dict(torch.load(ckpt_path, map_location='cpu')['state_dict'])\n",
    "model.load_state_dict(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([[ 1.9660, -1.8921]], device='cuda:3', grad_fn=<AddmmBackward0>)\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 DeepLIFT\n",
    "\n",
    "explainer = DeepLIFT(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\n",
    "\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 3421\n",
      "explain graph line 2536\n",
      "explain graph line 7266\n",
      "explain graph line 2650\n",
      "explain graph line 793\n",
      "explain graph line 7708\n",
      "explain graph line 5171\n",
      "explain graph line 5670\n",
      "explain graph line 1400\n",
      "explain graph line 655\n",
      "explain graph line 2746\n",
      "explain graph line 5111\n",
      "explain graph line 3959\n",
      "explain graph line 5602\n",
      "explain graph line 2661\n",
      "explain graph line 1500\n",
      "explain graph line 3117\n",
      "explain graph line 1056\n",
      "explain graph line 6898\n",
      "explain graph line 2455\n",
      "explain graph line 4479\n",
      "explain graph line 408\n",
      "explain graph line 892\n",
      "explain graph line 5199\n",
      "explain graph line 7001\n",
      "explain graph line 2840\n",
      "explain graph line 5358\n",
      "explain graph line 3284\n",
      "explain graph line 1390\n",
      "explain graph line 226\n",
      "explain graph line 7308\n",
      "explain graph line 6204\n",
      "explain graph line 194\n",
      "explain graph line 2501\n",
      "explain graph line 6398\n",
      "explain graph line 1191\n",
      "explain graph line 529\n",
      "explain graph line 6630\n",
      "explain graph line 7686\n",
      "explain graph line 6699\n",
      "explain graph line 969\n",
      "explain graph line 6030\n",
      "explain graph line 272\n",
      "explain graph line 1761\n",
      "explain graph line 5659\n",
      "explain graph line 745\n",
      "explain graph line 2783\n",
      "explain graph line 1170\n",
      "explain graph line 7711\n",
      "explain graph line 1668\n",
      "explain graph line 3061\n",
      "explain graph line 7670\n",
      "explain graph line 107\n",
      "explain graph line 2060\n",
      "explain graph line 7530\n",
      "explain graph line 3822\n",
      "explain graph line 7049\n",
      "explain graph line 4909\n",
      "explain graph line 6946\n",
      "explain graph line 1361\n",
      "explain graph line 6307\n",
      "explain graph line 7716\n",
      "explain graph line 7637\n",
      "explain graph line 6322\n",
      "explain graph line 7773\n",
      "explain graph line 6911\n",
      "explain graph line 4692\n",
      "explain graph line 331\n",
      "explain graph line 644\n",
      "explain graph line 1475\n",
      "explain graph line 1073\n",
      "explain graph line 5069\n",
      "explain graph line 7264\n",
      "explain graph line 2578\n",
      "explain graph line 4393\n",
      "explain graph line 2074\n",
      "explain graph line 7449\n",
      "explain graph line 3044\n",
      "explain graph line 4326\n",
      "explain graph line 4216\n",
      "explain graph line 4037\n",
      "explain graph line 5299\n",
      "explain graph line 1497\n",
      "explain graph line 7539\n",
      "explain graph line 5617\n",
      "explain graph line 1738\n",
      "explain graph line 1947\n",
      "explain graph line 1001\n",
      "explain graph line 4706\n",
      "explain graph line 2703\n",
      "explain graph line 4321\n",
      "explain graph line 6080\n",
      "explain graph line 5401\n",
      "explain graph line 481\n",
      "explain graph line 4676\n",
      "explain graph line 5934\n",
      "explain graph line 1535\n",
      "explain graph line 6181\n",
      "explain graph line 5968\n",
      "explain graph line 6690\n",
      "explain graph line 3776\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 = 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.1517\n",
      "Fidelity_inv: -0.0718\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
}