main_syn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f18e59c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import time\n",
    "from tqdm.notebook import tqdm\n",
    "import numpy as np\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.utils import dense_to_sparse, add_self_loops, remove_self_loops, k_hop_subgraph\n",
    "from torch_geometric.nn.conv.gcn_conv import gcn_norm\n",
    "\n",
    "from utils import *\n",
    "from model import *\n",
    "from explainer import *\n",
    "from main_distiller import *\n",
    "\n",
    "torch.manual_seed(123)\n",
    "torch.random.manual_seed(123)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb53eb3",
   "metadata": {},
   "source": [
    "# DnX and FastDnX \n",
    "\n",
    "\n",
    "Example explaination:\n",
    "\n",
    "Folder description:\n",
    "\n",
    "  * trained_gcn: contains the models that will be explained\n",
    "  * trained_distiller: contains the distilled models\n",
    "  * explanation: contains the explanations generated by our models\n",
    "\n",
    "Datasets: \n",
    "\n",
    "The nomenclature (syn1, syn2 ...., syn6) is used to facilitate experiments\n",
    "\n",
    "  * BA-HouseShapes (syn1), \n",
    "  * BA-Community (syn2), \n",
    "  * BA-Grids (syn3), \n",
    "  * Tree-Cycles (syn4), \n",
    "  * TreeGrids (syn5), \n",
    "  * BA-Bottle-Shaped (syn6) \n",
    "\n",
    "To run the DnX and Fast DnX explainers for the synthetic datasets just run the notebook `main_syn.ipynb`:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4e59206f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'syn1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a6ddd22d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = 'datasets/'\n",
    "out_path = 'explanation/'\n",
    "distiller_path = 'trained_distiller/'\n",
    "model_path = 'trained_gcn/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc683e54",
   "metadata": {},
   "source": [
    "## Loading dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cbac3327",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trained_gcn/syn1.pth.tar\n",
      "=> loading checkpoint 'trained_gcn/syn1.pth.tar'\n"
     ]
    }
   ],
   "source": [
    "# Load dataset and model to be explainer \n",
    "\n",
    "A_np, X_np = load_XA(dataset, datadir = dataset_path)\n",
    "num_nodes = X_np.shape[0]\n",
    "labels = load_labels(dataset, datadir = dataset_path)\n",
    "num_class = max(labels) + 1\n",
    "\n",
    "ckpt = load_ckpt(dataset,datadir = model_path)\n",
    "layer = 3\n",
    "\n",
    "A = torch.tensor(A_np, dtype=torch.float32).to(device)\n",
    "X = torch.tensor(X_np, dtype=torch.float32).to(device)\n",
    "\n",
    "if dataset =='syn2':\n",
    "    X = torch.concat([F.one_hot(torch.sum(A,1).type(torch.LongTensor)).type(torch.float32), X],1)\n",
    "else: \n",
    "    X = F.one_hot(torch.sum(A,1).type(torch.LongTensor)).type(torch.float32).to(device)\n",
    "\n",
    "input_dim = X.shape[1]\n",
    "\n",
    "edge_index,_ = dense_to_sparse(A)\n",
    "edge_index = edge_index.to(device)\n",
    "\n",
    "pred = ckpt[\"save_data\"][\"pred\"].squeeze(0)\n",
    "L_model = torch.softmax(torch.tensor(pred),1)\n",
    "\n",
    "\n",
    "node_list, k = get_nodes_explained(dataset, A_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f60d62e",
   "metadata": {},
   "source": [
    "## Generating the distilled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b7f03c5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9844db2b9b954fdcbe927ad83bb29773",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/999 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean Accuracy:0.9267857142857143 |  Std Accuracy:0.0\n",
      "trained_distiller/SGC_syn1.pth.tar\n"
     ]
    }
   ],
   "source": [
    "main(dataset, (X,  A, edge_index, L_model ), 1000, 1, layer, input_dim,num_class, distiller_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b8bd1489",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trained_distiller/SGC_syn1.pth.tar\n",
      "=> loading checkpoint 'trained_distiller/SGC_syn1.pth.tar'\n"
     ]
    }
   ],
   "source": [
    "ckpt_distillation = load_ckpt('SGC_'+dataset,datadir = distiller_path)\n",
    "\n",
    "model = SGC(3, input_dim, num_class).to(device)\n",
    "model.load_state_dict(ckpt_distillation[\"model_state\"])\n",
    "model.eval()\n",
    "pred_model = model(X,edge_index).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5539f58",
   "metadata": {},
   "source": [
    "# Generating explanations for the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c9751e",
   "metadata": {},
   "source": [
    "## DnX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fd6bfcd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Explaining syn1 dataset \n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "557bc098b4bc4713a9db941700c6fc94",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/400 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy:  0.975\n",
      "Precision:  0.975\n"
     ]
    }
   ],
   "source": [
    "nodes_explanations = {}\n",
    "results = {}\n",
    "\n",
    "t = 5\n",
    "itr_no = 0\n",
    "\n",
    "print(\"Explaining {} dataset \".format(dataset))\n",
    "\n",
    "for node in tqdm(node_list):\n",
    "    nodes_explanations_aux = {}\n",
    "    acc_top = 0\n",
    "\n",
    "    neighbors, sub_edge_index, node_idx_new, _ = k_hop_subgraph(int(node), layer, edge_index, relabel_nodes=True)\n",
    "    sub_X = X[neighbors]\n",
    "\n",
    "    node_idx_new, sub_edge_index, sub_X, neighbors = node_idx_new.to(device), sub_edge_index.to(device), sub_X.to(\n",
    "        device), neighbors.to(device)\n",
    "    \n",
    "    explainer = DnX(len(neighbors), node_idx_new).to(device)\n",
    "    opt = torch.optim.Adam(explainer.parameters(), lr=0.1)\n",
    "\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode='min',\n",
    "                                                           factor=0.5, min_lr=1e-5,\n",
    "                                                           patience=20,\n",
    "                                                           verbose=True)\n",
    "    loss = nn.MSELoss()\n",
    "\n",
    "    for epoch in range(100):\n",
    "\n",
    "        explainer.train()\n",
    "        opt.zero_grad()\n",
    "\n",
    "        expl, pred_ex = explainer(sub_X, sub_edge_index, model, t)\n",
    "\n",
    "        l = loss(pred_ex[node_idx_new], pred_model[node])\n",
    "\n",
    "        l.backward(retain_graph=True)\n",
    "\n",
    "        opt.step()\n",
    "        scheduler.step(l)\n",
    "\n",
    "        nodes_explanations_aux[node] = generating_explanations(node, model, explainer, edge_index, X, t, k)[node]\n",
    "        acc, prec = evaluate_syn_explanation(nodes_explanations_aux, dataset)\n",
    "        if acc > acc_top: \n",
    "            acc_top = acc\n",
    "            exp_top = nodes_explanations_aux[node]\n",
    "\n",
    "        if acc == 1.0:\n",
    "            break\n",
    "\n",
    "    nodes_explanations[node] = exp_top\n",
    "    acc, prec = evaluate_syn_explanation(nodes_explanations, dataset)\n",
    "\n",
    "\n",
    "with open(out_path+'nodes_explanations_DnX_{}.txt'.format(dataset), 'w') as f:\n",
    "    f.write(\"%s\\n\" % nodes_explanations)\n",
    "\n",
    "print(\"Accuracy: \", acc)\n",
    "print(\"Precision: \", prec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6a46f7c",
   "metadata": {},
   "source": [
    "## FastDnX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "228dcc64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Explaining syn1 dataset \n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdc2d4a7d391484da6f7714a57941570",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/400 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy:  0.995\n",
      "Precision:  0.995\n"
     ]
    }
   ],
   "source": [
    "print(\"Explaining {} dataset \".format(dataset))\n",
    "\n",
    "explanations = {}\n",
    "\n",
    "W = ckpt_distillation['model_state']['conv.lin.weight'].to(device)\n",
    "bias = (ckpt_distillation['model_state']['conv.lin.bias']).to(device)\n",
    "A_pot = A_k_hop(A, layer).to(device)\n",
    "\n",
    "k_nodes = k\n",
    "\n",
    "for no_alvo in tqdm(node_list):\n",
    "        nodes_neigh, _, node_ex, _ = k_hop_subgraph(int(no_alvo), layer, edge_index)\n",
    "\n",
    "        S = (X[nodes_neigh].T * A_pot[no_alvo, nodes_neigh]).T       \n",
    "        pred = torch.matmul(S, W.T)   \n",
    "        L = torch.ones(len(nodes_neigh), len(pred_model[no_alvo])).to(device) * pred_model[no_alvo] - bias#pred_model[no_alvo] - bias \n",
    "        expl = torch.diag(torch.matmul(L, pred.T) )\n",
    "\n",
    "        if len(nodes_neigh)<k:\n",
    "            k_nodes= len(nodes_neigh)\n",
    "        values, nodes = torch.topk(expl, dim=0,k=k_nodes)\n",
    "        k_nodes = k\n",
    "\n",
    "        explanations[no_alvo] = nodes_neigh[nodes].tolist()\n",
    "    \n",
    "acc, prec = evaluate_syn_explanation(explanations,dataset)\n",
    "\n",
    "print(\"Accuracy: \", acc)\n",
    "print(\"Precision: \", prec)\n",
    "\n",
    "with open(out_path+'nodes_explanations_fastDnX_{}.txt'.format(dataset), 'w') as f:\n",
    "    f.write(\"%s\\n\" % explanations)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14cffb0c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}