main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "11b50bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from  dig import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "932d5c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "import torch_geometric\n",
    "from torch_geometric.data import Dataset, Data\n",
    "import numpy as np \n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import torch.nn.functional as F \n",
    "from torch.nn import Linear, BatchNorm1d, ModuleList\n",
    "from torch_geometric.nn import TransformerConv, TopKPooling \n",
    "from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp\n",
    "from torch_geometric.data import DataLoader\n",
    "import torch\n",
    "import torch.nn.functional as F \n",
    "from torch.nn import Sequential, Linear, BatchNorm1d, ReLU\n",
    "from torch_geometric.nn import TransformerConv, GATConv, TopKPooling, BatchNorm\n",
    "from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp\n",
    "from torch_geometric.nn.conv.x_conv import XConv\n",
    "import torch \n",
    "from torch_geometric.data import DataLoader\n",
    "from sklearn.metrics import confusion_matrix, f1_score, \\\n",
    "    accuracy_score, precision_score, recall_score, roc_auc_score\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6ab461a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Torch version: 1.11.0\n",
      "Cuda available: True\n",
      "Torch geometric version: 2.0.5.dev20220428\n"
     ]
    }
   ],
   "source": [
    "print(f\"Torch version: {torch.__version__}\")\n",
    "print(f\"Cuda available: {torch.cuda.is_available()}\")\n",
    "print(f\"Torch geometric version: {torch_geometric.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "35bff41a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "path_temp= \"dataset-call-graph-blogpost-material/dataset/\"\n",
    "good_data=pickle.load(open(path_temp+\"goodware_graphs.p\",\"rb\"))\n",
    "bad_data=pickle.load(open(path_temp+\"malware_graphs.p\",\"rb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d4f8dd",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "16078f68",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MoleculeDataset(Dataset):\n",
    "    def __init__(self, root, filename,good_data, bad_data, test=False, transform=None, pre_transform=None):\n",
    "        \"\"\"\n",
    "        root = Where the dataset should be stored. This folder is split\n",
    "        into raw_dir (downloaded dataset) and processed_dir (processed data). \n",
    "        \"\"\"\n",
    "        self.test = test\n",
    "        self.filename = filename\n",
    "        self.good_data=good_data\n",
    "        self.bad_data=bad_data\n",
    "        super(MoleculeDataset, self).__init__(root, transform, pre_transform)\n",
    "        self.data_passed=self.good_data+self.bad_data\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        \"\"\" If this file exists in raw_dir, the download is not triggered.\n",
    "            (The download func. is not implemented here)  \n",
    "        \"\"\"\n",
    "        return [\"goodware_graphs.p\",\"malware_graphs.p\"]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        \"\"\" If these files are found in raw_dir, processing is skipped\"\"\"\n",
    "#         self.data = self.data_passed\n",
    "        return \"NOt implemented yet\"\n",
    "\n",
    "        # if self.test:\n",
    "        #     return [f'data_test_{i}.pt' for i in list(self.data.index)]\n",
    "        # else:\n",
    "        #     return [f'data_{i}.pt' for i in list(self.data.index)]\n",
    "\n",
    "    def download(self):\n",
    "        pass\n",
    "\n",
    "    def process(self):\n",
    "#         self.data = self.data_passed\n",
    "        print(\"started doing stuff\")\n",
    "        index=0\n",
    "\n",
    "        for mol in (self.good_data):\n",
    "            \n",
    "\n",
    "            feature_data=mol[0]\n",
    "            neighbour_data=mol[1]\n",
    "\n",
    "            # i need to create inst ->index\n",
    "            instr_index=dict()\n",
    "            for a,b in enumerate(feature_data):\n",
    "              instr_index[b]=a \n",
    "\n",
    "            # Get node features\n",
    "            node_feats = self._get_node_features(feature_data,instr_index)\n",
    "            # Get edge features\n",
    "            edge_feats = self._get_edge_features(neighbour_data)\n",
    "#             edge_feats = []\n",
    "            # Get adjacency info\n",
    "            edge_index = self._get_adjacency_info(neighbour_data,instr_index)\n",
    "\n",
    "            # Get labels info\n",
    "            label = 1\n",
    "\n",
    "            # Create data object\n",
    "            data = Data(x=node_feats, \n",
    "                        edge_index=edge_index,\n",
    "                        edge_attr=edge_feats,\n",
    "                        y=label,\n",
    "                        smiles=mol\n",
    "                        ) \n",
    "                        \n",
    "            torch.save(data, \n",
    "                    os.path.join(self.processed_dir, \n",
    "                                 f'data_{index}.pt'))\n",
    "\n",
    "            index+=1\n",
    "        \n",
    "        for mol in (self.bad_data):\n",
    "            \n",
    "\n",
    "            feature_data=mol[0]\n",
    "            neighbour_data=mol[1]\n",
    "\n",
    "            # i need to create inst ->index\n",
    "            instr_index=dict()\n",
    "            for a,b in enumerate(feature_data):\n",
    "              instr_index[b]=a \n",
    "\n",
    "            # Get node features\n",
    "            node_feats = self._get_node_features(feature_data,instr_index)\n",
    "            # Get edge features\n",
    "            edge_feats = self._get_edge_features(neighbour_data)\n",
    "#             edge_feats = []\n",
    "            # Get adjacency info\n",
    "            edge_index = self._get_adjacency_info(neighbour_data,instr_index)\n",
    "\n",
    "            # Get labels info\n",
    "            label = 0\n",
    "\n",
    "            # Create data object\n",
    "            data = Data(x=node_feats, \n",
    "                        edge_index=edge_index,\n",
    "                        edge_attr=edge_feats,\n",
    "                        y=label,\n",
    "                        smiles=mol\n",
    "                        ) \n",
    "                        \n",
    "            torch.save(data, \n",
    "                    os.path.join(self.processed_dir, \n",
    "                                 f'data_{index}.pt'))\n",
    "\n",
    "            index+=1\n",
    "        \n",
    "        \n",
    "\n",
    "    def process_node_features(self,val, node_feature_menu):\n",
    "      # print(\"jj\")\n",
    "      # print(val)\n",
    "      node_out=[]\n",
    "      keys=val.keys()\n",
    "      for name in node_feature_menu : \n",
    "        if(name in keys):\n",
    "          node_out.append(val[name])\n",
    "        else:\n",
    "          node_out.append(0)\n",
    "      return node_out\n",
    "\n",
    "\n",
    "    def _get_node_features(self, node_features,instr_index):\n",
    "        \"\"\" \n",
    "        This will return a matrix / 2d array of the shape\n",
    "        [Number of Nodes, Node Feature size]\n",
    "        \"\"\"\n",
    "        node_feature_menu=['mov', 'call', 'lea', 'jmp', 'push', 'add', 'xor', 'cmp', 'int3', 'nop', 'pushl', 'dec', 'sub', 'insl', 'inc','jz', 'jnz', 'je', 'jne', 'ja', 'jna', 'js', 'jns', 'jl', 'jnl', 'jg', 'jng']\n",
    "\n",
    "        all_node_feats = []\n",
    "\n",
    "        for atom in node_features:\n",
    "            # print(\"atom is\",atom)\n",
    "            \n",
    "            node_feats = []\n",
    "            node_feats = self.process_node_features(node_features[atom],node_feature_menu)\n",
    "            # Append node features to matrix\n",
    "            all_node_feats.append(node_feats)\n",
    "\n",
    "        all_node_feats = np.asarray(all_node_feats)\n",
    "        # all_node_feats = np.transpose(all_node_feats)\n",
    "        return torch.tensor(all_node_feats, dtype=torch.float)\n",
    "\n",
    "    def _get_edge_features(self, mol):\n",
    "        \"\"\" \n",
    "        This will return a matrix / 2d array of the shape\n",
    "        [Number of edges, Edge Feature size]\n",
    "        \"\"\"\n",
    "        all_edge_feats = []\n",
    "\n",
    "        # for bond in mol.GetBonds():\n",
    "        #     edge_feats = []\n",
    "        #     # Feature 1: Bond type (as double)\n",
    "        #     edge_feats.append(bond.GetBondTypeAsDouble())\n",
    "        #     # Feature 2: Rings\n",
    "        #     edge_feats.append(bond.IsInRing())\n",
    "        #     # Append node features to matrix (twice, per direction)\n",
    "        #     all_edge_feats += [edge_feats, edge_feats]\n",
    "\n",
    "        all_edge_feats = np.asarray(all_edge_feats)\n",
    "        return torch.tensor(all_edge_feats, dtype=torch.float)\n",
    "\n",
    "    def get_one_instr_adjacency(self,val,instr_index,self_instruction_number):\n",
    "        edge_val=[]\n",
    "        for a in val:\n",
    "            edge_val+=[[instr_index[self_instruction_number],instr_index[a]]]\n",
    "\n",
    "        return edge_val\n",
    "\n",
    "    def _get_adjacency_info(self, mol, instr_index):\n",
    "        \"\"\"\n",
    "        We could also use rdmolops.GetAdjacencyMatrix(mol)\n",
    "        but we want to be sure that the order of the indices\n",
    "        matches the order of the edge features\n",
    "        \"\"\"\n",
    "        edge_indices = []\n",
    "        for bond in mol:\n",
    "            edge_indices += self.get_one_instr_adjacency(mol[bond],instr_index,bond)\n",
    "\n",
    "        # print(\"--------\")\n",
    "        # print(len(edge_indices))\n",
    "        # print(len(edge_indices[0]))\n",
    "        # print(len(edge_indices[1]))\n",
    "        # print(edge_indices)\n",
    "        # edge_indices = torch.tensor(edge_indices)\n",
    "        # edge_indices = edge_indices.t().to(torch.long).view(2, -1)\n",
    "\n",
    "        np_edge=np.array(edge_indices)\n",
    "        np_edge=np.transpose(np_edge)\n",
    "        tr_edge=torch.tensor(np_edge, dtype=torch.long)\n",
    "\n",
    "\n",
    "        return tr_edge\n",
    "\n",
    "    def _get_labels(self, label):\n",
    "        label = np.asarray([label])\n",
    "        return torch.tensor(label, dtype=torch.int64)\n",
    "\n",
    "    def len(self):\n",
    "#         return self.data_pass.shape[0]\n",
    "        return len(self.data_passed)\n",
    "\n",
    "    def get(self, idx):\n",
    "        \"\"\" - Equivalent to __getitem__ in pytorch\n",
    "            - Is not needed for PyG's InMemoryDataset\n",
    "        \"\"\"\n",
    "        data = torch.load(os.path.join(self.processed_dir, \n",
    "                             f'data_{idx}.pt'))\n",
    "        return data\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9fbe3557",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "started doing stuff\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# train_dataset = MoleculeDataset(root=\"data/\", filename=\"goodware_graphs.p\", good_data=good_data[:100], bad_data=bad_data[:100])\n",
    "train_dataset = MoleculeDataset(root=\"data/\", filename=\"goodware_graphs.p\", good_data=good_data, bad_data=bad_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b859ca93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # test_cut_dataset=train_dataset[526:566]\n",
    "# # train_cut_dataset= train_dataset[:526]+train_dataset[566:]\n",
    "\n",
    "\n",
    "# NUM_GRAPHS_PER_BATCH = 1\n",
    "\n",
    "# train_loader = DataLoader(train_dataset, \n",
    "#                     batch_size=NUM_GRAPHS_PER_BATCH,\n",
    "#                      shuffle=False,\n",
    "#                          num_workers=0)\n",
    "# test_loader = DataLoader(train_dataset, \n",
    "#                     batch_size=NUM_GRAPHS_PER_BATCH,\n",
    "#                      shuffle=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7db406",
   "metadata": {},
   "source": [
    "## Model recipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e58fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "FileName: models.py\n",
    "Description: GNN models' set\n",
    "Time: 2020/7/30 9:01\n",
    "Project: GNN_benchmark\n",
    "Author: Shurui Gui\n",
    "\"\"\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch_geometric.nn as gnn\n",
    "from torch_geometric.utils.loop import add_self_loops, remove_self_loops\n",
    "from torch_geometric.data.batch import Batch\n",
    "\n",
    "from typing import Callable, Union, Tuple\n",
    "from torch_geometric.typing import OptPairTensor, Adj, OptTensor, Size\n",
    "from torch import Tensor\n",
    "\n",
    "from torch_sparse import SparseTensor\n",
    "\n",
    "\n",
    "class GNNBasic(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def arguments_read(self, *args, **kwargs):\n",
    "\n",
    "        data: Batch = kwargs.get('data') or None\n",
    "\n",
    "        if not data:\n",
    "            if not args:\n",
    "                assert 'x' in kwargs\n",
    "                assert 'edge_index' in kwargs\n",
    "                x, edge_index = kwargs['x'], kwargs['edge_index'],\n",
    "                batch = kwargs.get('batch')\n",
    "                if batch is None:\n",
    "                    batch = torch.zeros(kwargs['x'].shape[0], dtype=torch.int64, device=x.device)\n",
    "            elif len(args) == 2:\n",
    "                x, edge_index = args[0], args[1]\n",
    "                batch = torch.zeros(args[0].shape[0], dtype=torch.int64, device=x.device)\n",
    "            elif len(args) == 3:\n",
    "                x, edge_index, batch = args[0], args[1], args[2]\n",
    "            else:\n",
    "                raise ValueError(f\"forward's args should take 2 or 3 arguments but got {len(args)}\")\n",
    "        else:\n",
    "            x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "\n",
    "        return x, edge_index, batch\n",
    "\n",
    "\n",
    "class GCN_3l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 3\n",
    "\n",
    "        self.conv1 = GCNConv(dim_node, dim_hidden)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNConv(dim_hidden, dim_hidden)\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GCN_2l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GCNConv(dim_node, dim_hidden)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNConv(dim_hidden, dim_hidden)\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        \n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "            \n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GIN_3l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 3\n",
    "\n",
    "        self.conv1 = GINConv(nn.Sequential(nn.Linear(dim_node, dim_hidden), nn.ReLU(),\n",
    "                                           nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                           # nn.BatchNorm1d(dim_hidden)))\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GINConv(nn.Sequential(nn.Linear(dim_hidden, dim_hidden), nn.ReLU(),\n",
    "                                      nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                      # nn.BatchNorm1d(dim_hidden)))\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GIN_2l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GINConv(nn.Sequential(nn.Linear(dim_node, dim_hidden), nn.ReLU(),\n",
    "                                           nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                           # nn.BatchNorm1d(dim_hidden)))\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GINConv(nn.Sequential(nn.Linear(dim_hidden, dim_hidden), nn.ReLU(),\n",
    "                                      nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                      # nn.BatchNorm1d(dim_hidden)))\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GCNConv(gnn.GCNConv):\n",
    "\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.edge_weight = None\n",
    "\n",
    "    def forward(self, x: Tensor, edge_index: Adj,\n",
    "                edge_weight: OptTensor = None) -> Tensor:\n",
    "        \"\"\"\"\"\"\n",
    "\n",
    "        if self.normalize and edge_weight is None:\n",
    "            if isinstance(edge_index, Tensor):\n",
    "                cache = self._cached_edge_index\n",
    "                if cache is None:\n",
    "                    edge_index, edge_weight = gnn.conv.gcn_conv.gcn_norm(  # yapf: disable\n",
    "                        edge_index, edge_weight, x.size(self.node_dim),\n",
    "                        self.improved, self.add_self_loops, dtype=x.dtype)\n",
    "                    if self.cached:\n",
    "                        self._cached_edge_index = (edge_index, edge_weight)\n",
    "                else:\n",
    "                    edge_index, edge_weight = cache[0], cache[1]\n",
    "\n",
    "            elif isinstance(edge_index, SparseTensor):\n",
    "                cache = self._cached_adj_t\n",
    "                if cache is None:\n",
    "                    edge_index = gnn.conv.gcn_conv.gcn_norm(  # yapf: disable\n",
    "                        edge_index, edge_weight, x.size(self.node_dim),\n",
    "                        self.improved, self.add_self_loops, dtype=x.dtype)\n",
    "                    if self.cached:\n",
    "                        self._cached_adj_t = edge_index\n",
    "                else:\n",
    "                    edge_index = cache\n",
    "\n",
    "        # --- add require_grad ---\n",
    "        edge_weight.requires_grad_(True)\n",
    "\n",
    "        x = torch.matmul(x, self.weight)\n",
    "\n",
    "        # propagate_type: (x: Tensor, edge_weight: OptTensor)\n",
    "        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,\n",
    "                             size=None)\n",
    "\n",
    "        if self.bias is not None:\n",
    "            out += self.bias\n",
    "\n",
    "        # --- My: record edge_weight ---\n",
    "        self.edge_weight = edge_weight\n",
    "\n",
    "        return out\n",
    "\n",
    "    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):\n",
    "        size = self.__check_input__(edge_index, size)\n",
    "\n",
    "        # Run \"fused\" message and aggregation (if applicable).\n",
    "        if (isinstance(edge_index, SparseTensor) and self.fuse\n",
    "                and not self.__explain__):\n",
    "            coll_dict = self.__collect__(self.__fused_user_args__, edge_index,\n",
    "                                         size, kwargs)\n",
    "\n",
    "            msg_aggr_kwargs = self.inspector.distribute(\n",
    "                'message_and_aggregate', coll_dict)\n",
    "            out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)\n",
    "\n",
    "            update_kwargs = self.inspector.distribute('update', coll_dict)\n",
    "            return self.update(out, **update_kwargs)\n",
    "\n",
    "        # Otherwise, run both functions in separation.\n",
    "        elif isinstance(edge_index, Tensor) or not self.fuse:\n",
    "            coll_dict = self.__collect__(self.__user_args__, edge_index, size,\n",
    "                                         kwargs)\n",
    "\n",
    "            msg_kwargs = self.inspector.distribute('message', coll_dict)\n",
    "            out = self.message(**msg_kwargs)\n",
    "\n",
    "            # For `GNNExplainer`, we require a separate message and aggregate\n",
    "            # procedure since this allows us to inject the `edge_mask` into the\n",
    "            # message passing computation scheme.\n",
    "            if self.__explain__:\n",
    "                edge_mask = self.__edge_mask__\n",
    "                # Some ops add self-loops to `edge_index`. We need to do the\n",
    "                # same for `edge_mask` (but do not train those).\n",
    "                if out.size(self.node_dim) != edge_mask.size(0):\n",
    "                    loop = edge_mask.new_ones(size[0])\n",
    "                    edge_mask = torch.cat([edge_mask, loop], dim=0)\n",
    "                assert out.size(self.node_dim) == edge_mask.size(0)\n",
    "                out = out * edge_mask.view([-1] + [1] * (out.dim() - 1))\n",
    "\n",
    "            aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)\n",
    "            out = self.aggregate(out, **aggr_kwargs)\n",
    "\n",
    "            update_kwargs = self.inspector.distribute('update', coll_dict)\n",
    "            return self.update(out, **update_kwargs)\n",
    "\n",
    "\n",
    "class GINConv(gnn.GINConv):\n",
    "\n",
    "    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,\n",
    "                 **kwargs):\n",
    "        super().__init__(nn, eps, train_eps, **kwargs)\n",
    "        self.edge_weight = None\n",
    "        self.fc_steps = None\n",
    "        self.reweight = None\n",
    "\n",
    "    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,\n",
    "                edge_weight: OptTensor = None, task='explain', **kwargs) -> Tensor:\n",
    "        \"\"\"\"\"\"\n",
    "        self.num_nodes = x.shape[0]\n",
    "        if isinstance(x, Tensor):\n",
    "            x: OptPairTensor = (x, x)\n",
    "\n",
    "        # propagate_type: (x: OptPairTensor)\n",
    "        if edge_weight is not None:\n",
    "            self.edge_weight = edge_weight\n",
    "            assert edge_weight.shape[0] == edge_index.shape[1]\n",
    "            self.reweight = False\n",
    "        else:\n",
    "            edge_index, _ = remove_self_loops(edge_index)\n",
    "            self_loop_edge_index, _ = add_self_loops(edge_index, num_nodes=self.num_nodes)\n",
    "            if self_loop_edge_index.shape[1] != edge_index.shape[1]:\n",
    "                edge_index = self_loop_edge_index\n",
    "            self.reweight = True\n",
    "        out = self.propagate(edge_index, x=x[0], size=None)\n",
    "\n",
    "        if task == 'explain':\n",
    "            layer_extractor = []\n",
    "            hooks = []\n",
    "\n",
    "            def register_hook(module: nn.Module):\n",
    "                if not list(module.children()):\n",
    "                    hooks.append(module.register_forward_hook(forward_hook))\n",
    "\n",
    "            def forward_hook(module: nn.Module, input: Tuple[Tensor], output: Tensor):\n",
    "                # input contains x and edge_index\n",
    "                layer_extractor.append((module, input[0], output))\n",
    "\n",
    "            # --- register hooks ---\n",
    "            self.nn.apply(register_hook)\n",
    "\n",
    "            nn_out = self.nn(out)\n",
    "\n",
    "            for hook in hooks:\n",
    "                hook.remove()\n",
    "\n",
    "            fc_steps = []\n",
    "            step = {'input': None, 'module': [], 'output': None}\n",
    "            for layer in layer_extractor:\n",
    "                if isinstance(layer[0], nn.Linear):\n",
    "                    if step['module']:\n",
    "                        fc_steps.append(step)\n",
    "                    # step = {'input': layer[1], 'module': [], 'output': None}\n",
    "                    step = {'input': None, 'module': [], 'output': None}\n",
    "                step['module'].append(layer[0])\n",
    "                if kwargs.get('probe'):\n",
    "                    step['output'] = layer[2]\n",
    "                else:\n",
    "                    step['output'] = None\n",
    "\n",
    "            if step['module']:\n",
    "                fc_steps.append(step)\n",
    "            self.fc_steps = fc_steps\n",
    "        else:\n",
    "            nn_out = self.nn(out)\n",
    "\n",
    "\n",
    "        return nn_out\n",
    "\n",
    "    def message(self, x_j: Tensor) -> Tensor:\n",
    "        if self.reweight:\n",
    "            edge_weight = torch.ones(x_j.shape[0], device=x_j.device)\n",
    "            edge_weight.data[-self.num_nodes:] += self.eps\n",
    "            edge_weight = edge_weight.detach().clone()\n",
    "            edge_weight.requires_grad_(True)\n",
    "            self.edge_weight = edge_weight\n",
    "        return x_j * self.edge_weight.view(-1, 1)\n",
    "\n",
    "    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):\n",
    "        size = self.__check_input__(edge_index, size)\n",
    "\n",
    "        # Run \"fused\" message and aggregation (if applicable).\n",
    "        if (isinstance(edge_index, SparseTensor) and self.fuse\n",
    "                and not self.__explain__):\n",
    "            coll_dict = self.__collect__(self.__fused_user_args__, edge_index,\n",
    "                                         size, kwargs)\n",
    "\n",
    "            msg_aggr_kwargs = self.inspector.distribute(\n",
    "                'message_and_aggregate', coll_dict)\n",
    "            out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)\n",
    "\n",
    "            update_kwargs = self.inspector.distribute('update', coll_dict)\n",
    "            return self.update(out, **update_kwargs)\n",
    "\n",
    "        # Otherwise, run both functions in separation.\n",
    "        elif isinstance(edge_index, Tensor) or not self.fuse:\n",
    "            coll_dict = self.__collect__(self.__user_args__, edge_index, size,\n",
    "                                         kwargs)\n",
    "\n",
    "            msg_kwargs = self.inspector.distribute('message', coll_dict)\n",
    "            out = self.message(**msg_kwargs)\n",
    "\n",
    "            # For `GNNExplainer`, we require a separate message and aggregate\n",
    "            # procedure since this allows us to inject the `edge_mask` into the\n",
    "            # message passing computation scheme.\n",
    "            if self.__explain__:\n",
    "                edge_mask = self.__edge_mask__\n",
    "                # Some ops add self-loops to `edge_index`. We need to do the\n",
    "                # same for `edge_mask` (but do not train those).\n",
    "                if out.size(self.node_dim) != edge_mask.size(0):\n",
    "                    loop = edge_mask.new_ones(size[0])\n",
    "                    edge_mask = torch.cat([edge_mask, loop], dim=0)\n",
    "                assert out.size(self.node_dim) == edge_mask.size(0)\n",
    "                out = out * edge_mask.view([-1] + [1] * (out.dim() - 1))\n",
    "\n",
    "            aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)\n",
    "            out = self.aggregate(out, **aggr_kwargs)\n",
    "\n",
    "            update_kwargs = self.inspector.distribute('update', coll_dict)\n",
    "            return self.update(out, **update_kwargs)\n",
    "\n",
    "\n",
    "class GNNPool(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "\n",
    "class GlobalMeanPool(GNNPool):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x, batch):\n",
    "        return gnn.global_mean_pool(x, batch)\n",
    "\n",
    "\n",
    "class IdenticalPool(GNNPool):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x, batch):\n",
    "        return x\n",
    "\n",
    "\n",
    "class GraphSequential(nn.Sequential):\n",
    "\n",
    "    def __init__(self, *args):\n",
    "        super().__init__(*args)\n",
    "\n",
    "    def forward(self, *input) -> Tensor:\n",
    "        for module in self:\n",
    "            if isinstance(input, tuple):\n",
    "                input = module(*input)\n",
    "            else:\n",
    "                input = module(input)\n",
    "        return input\n",
    "\n",
    "\n",
    "# explain_mask in propagation haven't pass sigmoid func\n",
    "class GCN_2l_mask(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GCNConv_mask(dim_node, dim_hidden)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNConv_mask(dim_hidden, dim_hidden)\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GIN_2l_mask(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GINConv_mask(nn.Sequential(nn.Linear(dim_node, dim_hidden), nn.ReLU(),\n",
    "                                           nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                           # nn.BatchNorm1d(dim_hidden)))\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GINConv_mask(nn.Sequential(nn.Linear(dim_hidden, dim_hidden), nn.ReLU(),\n",
    "                                      nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                      # nn.BatchNorm1d(dim_hidden)))\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36f81ca",
   "metadata": {},
   "source": [
    "### Dont lookup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "69ffcaeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "FileName: models.py\n",
    "Description: GNN models' set\n",
    "Time: 2020/7/30 9:01\n",
    "Project: GNN_benchmark\n",
    "Author: Shurui Gui\n",
    "\"\"\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch_geometric.nn as gnn\n",
    "from torch_geometric.utils.loop import add_self_loops, remove_self_loops\n",
    "from torch_geometric.data.batch import Batch\n",
    "\n",
    "from typing import Callable, Union, Tuple\n",
    "from torch_geometric.typing import OptPairTensor, Adj, OptTensor, Size\n",
    "from torch import Tensor\n",
    "\n",
    "from torch_sparse import SparseTensor\n",
    "\n",
    "\n",
    "class GNNBasic(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def arguments_read(self, *args, **kwargs):\n",
    "\n",
    "        data: Batch = kwargs.get('data') or None\n",
    "\n",
    "        if not data:\n",
    "            if not args:\n",
    "                assert 'x' in kwargs\n",
    "                assert 'edge_index' in kwargs\n",
    "                x, edge_index = kwargs['x'], kwargs['edge_index'],\n",
    "                batch = kwargs.get('batch')\n",
    "                if batch is None:\n",
    "                    batch = torch.zeros(kwargs['x'].shape[0], dtype=torch.int64, device=x.device)\n",
    "            elif len(args) == 2:\n",
    "                x, edge_index = args[0], args[1]\n",
    "                batch = torch.zeros(args[0].shape[0], dtype=torch.int64, device=x.device)\n",
    "            elif len(args) == 3:\n",
    "                x, edge_index, batch = args[0], args[1], args[2]\n",
    "            else:\n",
    "                raise ValueError(f\"forward's args should take 2 or 3 arguments but got {len(args)}\")\n",
    "        else:\n",
    "            x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "\n",
    "        return x, edge_index, batch\n",
    "\n",
    "\n",
    "class GCN_3l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 3\n",
    "\n",
    "        self.conv1 = GCNConv(dim_node, dim_hidden)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNConv(dim_hidden, dim_hidden)\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GCN_2l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GCNConv(dim_node, dim_hidden)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNConv(dim_hidden, dim_hidden)\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        \n",
    "        post_conv = self.relu1(self.conv1(x, edge_index))\n",
    "        for conv, relu in zip(self.convs, self.relus):\n",
    "            post_conv = relu(conv(post_conv, edge_index))\n",
    "            \n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GIN_3l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 3\n",
    "\n",
    "        self.conv1 = GINConv(nn.Sequential(nn.Linear(dim_node, dim_hidden), nn.ReLU(),\n",
    "                                           nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                           # nn.BatchNorm1d(dim_hidden)))\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GINConv(nn.Sequential(nn.Linear(dim_hidden, dim_hidden), nn.ReLU(),\n",
    "                                      nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                      # nn.BatchNorm1d(dim_hidden)))\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n",
    "\n",
    "\n",
    "class GIN_2l(GNNBasic):\n",
    "\n",
    "    def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n",
    "        super().__init__()\n",
    "        num_layer = 2\n",
    "\n",
    "        self.conv1 = GINConv(nn.Sequential(nn.Linear(dim_node, dim_hidden), nn.ReLU(),\n",
    "                                           nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                           # nn.BatchNorm1d(dim_hidden)))\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GINConv(nn.Sequential(nn.Linear(dim_hidden, dim_hidden), nn.ReLU(),\n",
    "                                      nn.Linear(dim_hidden, dim_hidden), nn.ReLU()))#,\n",
    "                                      # nn.BatchNorm1d(dim_hidden)))\n",
    "                for _ in range(num_layer - 1)\n",
    "             ]\n",
    "        )\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.relus = nn.ModuleList(\n",
    "            [\n",
    "                nn.ReLU()\n",
    "                for _ in range(num_layer - 1)\n",
    "            ]\n",
    "        )\n",
    "        if model_level == 'node':\n",
    "            self.readout = IdenticalPool()\n",
    "        else:\n",
    "            self.readout = GlobalMeanPool()\n",
    "\n",
    "        self.ffn = nn.Sequential(*(\n",
    "                [nn.Linear(dim_hidden, dim_hidden)] +\n",
    "                [nn.ReLU(), nn.Dropout(), nn.Linear(dim_hidden, num_classes)]\n",
    "        ))\n",
    "\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, *args, **kwargs) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        :param Required[data]: Batch - input data\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "\n",
    "\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "\n",
    "\n",
    "        out_readout = self.readout(post_conv, batch)\n",
    "\n",
    "        out = self.ffn(out_readout)\n",
    "        return out\n",
    "\n",
    "    def get_emb(self, *args, **kwargs) -> torch.Tensor:\n",
    "        x, edge_index, batch = self.arguments_read(*args, **kwargs)\n",
    "        post_conv = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            post_conv = conv(post_conv, edge_index)\n",
    "        return post_conv\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "962ad843",
   "metadata": {},
   "source": [
    "## Our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94d3a928",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "ccdbe958",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e31e6083",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim_req=train_dataset[0].x.shape[1]\n",
    "model = GCN_2l(model_level=\"graph_xx\", dim_node=dim_req, dim_hidden=dim_req*2,num_classes=2 ) \n",
    "    # def __init__(self, model_level, dim_node, dim_hidden, num_classes):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "10927db6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 4592\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "GCN_2l(\n",
       "  (conv1): GCNConv(27, 54)\n",
       "  (convs): ModuleList(\n",
       "    (0): GCNConv(54, 54)\n",
       "  )\n",
       "  (relu1): ReLU()\n",
       "  (relus): ModuleList(\n",
       "    (0): ReLU()\n",
       "  )\n",
       "  (readout): GlobalMeanPool()\n",
       "  (ffn): Sequential(\n",
       "    (0): Linear(in_features=54, out_features=2, bias=True)\n",
       "  )\n",
       "  (dropout): Dropout(p=0.5, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#%% Loading the model\n",
    "# model = GNN(feature_size=train_dataset[0].x.shape[0]) \n",
    "model = model.to(device)\n",
    "print(f\"Number of parameters: {count_parameters(model)}\")\n",
    "weights = torch.tensor([1, 1.2], dtype=torch.float32).to(device)\n",
    "loss_fn = torch.nn.CrossEntropyLoss(weight=weights)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.2, amsgrad=True)\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.08)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e97111",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "148a343d",
   "metadata": {},
   "source": [
    "# train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9689160d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch):\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    for batch in test_loader:\n",
    "        batch.to(device)  \n",
    "        pred = model(batch.x.float(), \n",
    "                        batch.edge_attr.float(),\n",
    "                        batch.edge_index, \n",
    "                        batch.batch\n",
    "                        ) \n",
    "        loss = torch.sqrt(loss_fn(pred, batch.y))    \n",
    "        all_preds.append(np.argmax(pred.cpu().detach().numpy(), axis=1))\n",
    "        all_labels.append(batch.y.cpu().detach().numpy())\n",
    "\n",
    "    all_preds = np.concatenate(all_preds).ravel()\n",
    "    all_labels = np.concatenate(all_labels).ravel()\n",
    "    calculate_metrics(all_preds, all_labels, epoch, \"test\")\n",
    "    return loss\n",
    "\n",
    "\n",
    "def calculate_metrics(y_pred, y_true, epoch, type):\n",
    "    print(f\"\\n Confusion matrix: \\n {confusion_matrix(y_pred, y_true)}\")\n",
    "    print(f\"F1 Score: {f1_score(y_pred, y_true)}\")\n",
    "    print(f\"Accuracy: {accuracy_score(y_pred, y_true)}\")\n",
    "    print(f\"Precision: {precision_score(y_pred, y_true)}\")\n",
    "    print(f\"Recall: {recall_score(y_pred, y_true)}\")\n",
    "    try:\n",
    "        roc = roc_auc_score(y_pred, y_true)\n",
    "        print(f\"ROC AUC: {roc}\")\n",
    "        # mlflow.log_metric(key=f\"ROC-AUC-{type}\", value=float(roc), step=epoch)\n",
    "    except:\n",
    "        # mlflow.log_metric(key=f\"ROC-AUC-{type}\", value=float(0), step=epoch)\n",
    "        print(f\"ROC AUC: notdefined\")\n",
    "        \n",
    "    # try:\n",
    "    #     # mlflow.log_metric(key=f\"accuracy-{type}\", value=float(accuracy_score(y_pred, y_true)), step=epoch)\n",
    "    # except:\n",
    "    #     # mlflow.log_metric(key=f\"ROC-AUC-{type}\", value=float(0), step=epoch)\n",
    "    #     print(f\"Accuracy: notdefined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "e09eaccc",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train(epoch):\n",
    "    # Enumerate over the data\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    for _, batch in enumerate(tqdm(train_loader)):\n",
    "        # Use GPU\n",
    "        batch.to(device)  \n",
    "        # Reset gradients\n",
    "\n",
    "        optimizer.zero_grad() \n",
    "        # Passing the node features and the connection info\n",
    "        pred = model(batch.x.float(), \n",
    "                                # batch.edge_attr.float(),\n",
    "                                batch.edge_index, \n",
    "                                batch.batch\n",
    "                                ) \n",
    "        # Calculating the loss and gradients\n",
    "#         loss = torch.sqrt(loss_fn(pred, batch.y)) \n",
    "        loss_v1 = loss_fn(pred, batch.y) \n",
    "        loss_v2=torch.sqrt(loss_v1)\n",
    "        \n",
    "        loss=loss_v2\n",
    "        loss.backward()  \n",
    "        \n",
    "        # this is the additional clipping\n",
    "        torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=7.0)\n",
    "        \n",
    "        \n",
    "        # Update using the gradients\n",
    "        optimizer.step()  \n",
    "\n",
    "\n",
    "        all_preds.append(np.argmax(pred.cpu().detach().numpy(), axis=1))\n",
    "        all_labels.append(batch.y.cpu().detach().numpy())\n",
    "    all_preds = np.concatenate(all_preds).ravel()\n",
    "    all_labels = np.concatenate(all_labels).ravel()\n",
    "    calculate_metrics(all_preds, all_labels, epoch, \"train\")\n",
    "    return loss\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "036bec8f",
   "metadata": {},
   "source": [
    "# lets train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "3a7f9c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import TransformerConv, GATConv, TopKPooling, BatchNorm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "ac796c4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1361 [00:35<?, ?it/s]\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'GCNConv' object has no attribute 'weight'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32md:\\code\\OELP_bigger\\main.ipynb Cell 24'\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=0'>1</a>\u001b[0m \u001b[39mfor\u001b[39;00m epoch \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39m20\u001b[39m):\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=1'>2</a>\u001b[0m     \u001b[39m# Training\u001b[39;00m\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=2'>3</a>\u001b[0m     model\u001b[39m.\u001b[39mtrain()\n\u001b[1;32m----> <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=3'>4</a>\u001b[0m     loss \u001b[39m=\u001b[39m train(epoch\u001b[39m=\u001b[39;49mepoch)\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=4'>5</a>\u001b[0m     loss \u001b[39m=\u001b[39m loss\u001b[39m.\u001b[39mdetach()\u001b[39m.\u001b[39mgpu()\u001b[39m.\u001b[39mnumpy()\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000016?line=5'>6</a>\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mEpoch \u001b[39m\u001b[39m{\u001b[39;00mepoch\u001b[39m}\u001b[39;00m\u001b[39m | Train Loss \u001b[39m\u001b[39m{\u001b[39;00mloss\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n",
      "\u001b[1;32md:\\code\\OELP_bigger\\main.ipynb Cell 21'\u001b[0m in \u001b[0;36mtrain\u001b[1;34m(epoch)\u001b[0m\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=9'>10</a>\u001b[0m         optimizer\u001b[39m.\u001b[39mzero_grad() \n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=10'>11</a>\u001b[0m         \u001b[39m# Passing the node features and the connection info\u001b[39;00m\n\u001b[1;32m---> <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=11'>12</a>\u001b[0m         pred \u001b[39m=\u001b[39m model(batch\u001b[39m.\u001b[39;49mx\u001b[39m.\u001b[39;49mfloat(), \n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=12'>13</a>\u001b[0m                                 \u001b[39m# batch.edge_attr.float(),\u001b[39;49;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=13'>14</a>\u001b[0m                                 batch\u001b[39m.\u001b[39;49medge_index, \n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=14'>15</a>\u001b[0m                                 batch\u001b[39m.\u001b[39;49mbatch\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=15'>16</a>\u001b[0m                                 ) \n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=16'>17</a>\u001b[0m         \u001b[39m# Calculating the loss and gradients\u001b[39;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=17'>18</a>\u001b[0m \u001b[39m#         loss = torch.sqrt(loss_fn(pred, batch.y)) \u001b[39;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000013?line=18'>19</a>\u001b[0m         loss_v1 \u001b[39m=\u001b[39m loss_fn(pred, batch\u001b[39m.\u001b[39my) \n",
      "File \u001b[1;32md:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1105'>1106</a>\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1106'>1107</a>\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1107'>1108</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1108'>1109</a>\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1109'>1110</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39m\u001b[39minput\u001b[39m, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1110'>1111</a>\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1111'>1112</a>\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "\u001b[1;32md:\\code\\OELP_bigger\\main.ipynb Cell 12'\u001b[0m in \u001b[0;36mGCN_2l.forward\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=137'>138</a>\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=138'>139</a>\u001b[0m \u001b[39m:param Required[data]: Batch - input data\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=139'>140</a>\u001b[0m \u001b[39m:return:\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=140'>141</a>\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=141'>142</a>\u001b[0m x, edge_index, batch \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39marguments_read(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m--> <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=143'>144</a>\u001b[0m post_conv \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mrelu1(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconv1(x, edge_index))\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=144'>145</a>\u001b[0m \u001b[39mfor\u001b[39;00m conv, relu \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconvs, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mrelus):\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=145'>146</a>\u001b[0m     post_conv \u001b[39m=\u001b[39m relu(conv(post_conv, edge_index))\n",
      "File \u001b[1;32md:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1105'>1106</a>\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1106'>1107</a>\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1107'>1108</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1108'>1109</a>\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1109'>1110</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39m\u001b[39minput\u001b[39m, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1110'>1111</a>\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1111'>1112</a>\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "\u001b[1;32md:\\code\\OELP_bigger\\main.ipynb Cell 12'\u001b[0m in \u001b[0;36mGCNConv.forward\u001b[1;34m(self, x, edge_index, edge_weight)\u001b[0m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=320'>321</a>\u001b[0m \u001b[39m# --- add require_grad ---\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=321'>322</a>\u001b[0m edge_weight\u001b[39m.\u001b[39mrequires_grad_(\u001b[39mTrue\u001b[39;00m)\n\u001b[1;32m--> <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=323'>324</a>\u001b[0m x \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mmatmul(x, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight)\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=325'>326</a>\u001b[0m \u001b[39m# propagate_type: (x: Tensor, edge_weight: OptTensor)\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=326'>327</a>\u001b[0m out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpropagate(edge_index, x\u001b[39m=\u001b[39mx, edge_weight\u001b[39m=\u001b[39medge_weight,\n\u001b[0;32m    <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/main.ipynb#ch0000008?line=327'>328</a>\u001b[0m                      size\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m)\n",
      "File \u001b[1;32md:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:1185\u001b[0m, in \u001b[0;36mModule.__getattr__\u001b[1;34m(self, name)\u001b[0m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1182'>1183</a>\u001b[0m     \u001b[39mif\u001b[39;00m name \u001b[39min\u001b[39;00m modules:\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1183'>1184</a>\u001b[0m         \u001b[39mreturn\u001b[39;00m modules[name]\n\u001b[1;32m-> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1184'>1185</a>\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mAttributeError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m object has no attribute \u001b[39m\u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m.\u001b[39mformat(\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1185'>1186</a>\u001b[0m     \u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m, name))\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'GCNConv' object has no attribute 'weight'"
     ]
    }
   ],
   "source": [
    "\n",
    "for epoch in range(20):\n",
    "    # Training\n",
    "    model.train()\n",
    "    loss = train(epoch=epoch)\n",
    "    loss = loss.detach().gpu().numpy()\n",
    "    print(f\"Epoch {epoch} | Train Loss {loss}\")\n",
    "    # mlflow.log_metric(key=\"Train loss\", value=float(loss), step=epoch)\n",
    "\n",
    "    # Testing\n",
    "    model.eval()\n",
    "    if epoch % 10 == 0 and epoch !=0:\n",
    "        loss = test(epoch=epoch)\n",
    "        loss = loss.detach().gpu().numpy()\n",
    "        print(f\"Epoch {epoch} | Test Loss {loss}\")\n",
    "        # mlflow.log_metric(key=\"Test loss\", value=float(loss), step=epoch)\n",
    "\n",
    "    scheduler.step()\n",
    "print(\"Done.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c2c903",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([464, 27]) torch.Size([2, 821])\n",
      "torch.Size([492, 27]) torch.Size([2, 1217])\n",
      "torch.Size([485, 27]) torch.Size([2, 925])\n",
      "torch.Size([440, 27]) torch.Size([2, 792])\n",
      "torch.Size([465, 27]) torch.Size([2, 836])\n",
      "torch.Size([468, 27]) torch.Size([2, 855])\n",
      "torch.Size([440, 27]) torch.Size([2, 792])\n",
      "torch.Size([587, 27]) torch.Size([2, 1083])\n",
      "torch.Size([1137, 27]) torch.Size([2, 3479])\n",
      "torch.Size([478, 27]) torch.Size([2, 866])\n",
      "torch.Size([1024, 27]) torch.Size([2, 2837])\n",
      "torch.Size([481, 27]) torch.Size([2, 690])\n",
      "torch.Size([488, 27]) torch.Size([2, 928])\n",
      "torch.Size([322, 27]) torch.Size([2, 354])\n",
      "torch.Size([1122, 27]) torch.Size([2, 2241])\n",
      "torch.Size([559, 27]) torch.Size([2, 982])\n",
      "torch.Size([1643, 27]) torch.Size([2, 4396])\n",
      "torch.Size([568, 27]) torch.Size([2, 1731])\n",
      "torch.Size([463, 27]) torch.Size([2, 816])\n",
      "torch.Size([484, 27]) torch.Size([2, 925])\n"
     ]
    }
   ],
   "source": [
    "# # for i in range(len(train_dataset)):\n",
    "#     # if (len(train_dataset[i].x[1])) == 768:\n",
    "#         # print(i)\n",
    "\n",
    "# for i in range(20):\n",
    "#     print(train_dataset[i].x.size(),train_dataset[i].edge_index.size())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92973361",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022/04/28 23:57:54 WARNING mlflow.utils.environment: Encountered an unexpected error while inferring pip requirements (model URI: C:\\Users\\JERRYJ~1\\AppData\\Local\\Temp\\tmpajrx1zfm\\model\\data, flavor: pytorch), fall back to return ['torch==1.11.0', 'cloudpickle==2.0.0']. Set logging level to DEBUG to see the full traceback.\n",
      "D:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\_distutils_hack\\__init__.py:30: UserWarning: Setuptools is replacing distutils.\n",
      "  warnings.warn(\"Setuptools is replacing distutils.\")\n"
     ]
    }
   ],
   "source": [
    "from datetime import date\n",
    "from datetime import datetime\n",
    "\n",
    "# %% Save the model \n",
    "model_type_name=\"20epoch_gpu\"\n",
    "\n",
    "mlflow.pytorch.log_model(model, \"model\") \n",
    "mlflow.end_run()\n",
    "\n",
    "\n",
    "# SAVE MODEL\n",
    "#  pls create a folder named model_saved if not there\n",
    "now = datetime.now()\n",
    "dt_string = now.strftime(\"%d_%m_%Y %H_%M_%S\")\n",
    "PATH_to_save_model=\"model_saved/\"+\"model_\"+dt_string+\":\"+model_type_name+\".pt\"\n",
    "PATH_to_save_model=\"model_saved/model \"+dt_string+model_type_name+\".pt\"\n",
    "torch.save(model, PATH_to_save_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa76746",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'model_saved/model 28_04_2022 23_57_541epoch_cpu.pt'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PATH_to_save_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "70addd8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PATH_to_save_model='model_saved/model 25_04_2022 13_52_341epoch_ennexp.pt'\n",
    "# PATH_to_save_model='model_saved/model 25_04_2022 14_50_061epoch_ennexp.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "f4982651",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GNN(\n",
       "  (conv1): GATConv(27, 256, heads=3)\n",
       "  (head_transform1): Linear(in_features=768, out_features=256, bias=True)\n",
       "  (pool1): TopKPooling(256, ratio=0.8, multiplier=1.0)\n",
       "  (conv2): GATConv(256, 256, heads=3)\n",
       "  (head_transform2): Linear(in_features=768, out_features=256, bias=True)\n",
       "  (pool2): TopKPooling(256, ratio=0.5, multiplier=1.0)\n",
       "  (conv3): GATConv(256, 256, heads=3)\n",
       "  (head_transform3): Linear(in_features=768, out_features=256, bias=True)\n",
       "  (pool3): TopKPooling(256, ratio=0.2, multiplier=1.0)\n",
       "  (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "  (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "  (linear3): Linear(in_features=512, out_features=128, bias=True)\n",
       "  (linear4): Linear(in_features=128, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model=torch.load(PATH_to_save_model)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XAI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d487bbc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GNNExplainer\n",
    "# Initialize explainer\n",
    "explainer = GNNExplainer(model, epochs=200, return_type='log_prob')\n",
    "graph = train_dataset[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "05615959",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 1217])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.edge_index.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9d2f6595",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.zeros(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "e3f91660",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "forward() got multiple values for argument 'edge_attr'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32md:\\code\\OELP_bigger\\try_24_04 GNN model for explainer.ipynb Cell 28'\u001b[0m in \u001b[0;36m<cell line: 4>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/try_24_04%20GNN%20model%20for%20explainer.ipynb#ch0000034?line=0'>1</a>\u001b[0m \u001b[39m# node_feat_mask, edge_mask = explainer.explain_graph(graph.x, graph.edge_index, batch_index=1)\u001b[39;00m\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/try_24_04%20GNN%20model%20for%20explainer.ipynb#ch0000034?line=1'>2</a>\u001b[0m \u001b[39m# node_feat_mask, edge_mask = explainer.explain_graph(graph.x, graph.edge_index, batch=3)\u001b[39;00m\n\u001b[0;32m      <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/try_24_04%20GNN%20model%20for%20explainer.ipynb#ch0000034?line=2'>3</a>\u001b[0m \u001b[39m# node_feat_mask, edge_mask = explainaer.explain_graph(graph.x, graph.edge_index, edge_attr=2)\u001b[39;00m\n\u001b[1;32m----> <a href='vscode-notebook-cell:/d%3A/code/OELP_bigger/try_24_04%20GNN%20model%20for%20explainer.ipynb#ch0000034?line=3'>4</a>\u001b[0m node_feat_mask, edge_mask \u001b[39m=\u001b[39m explainer\u001b[39m.\u001b[39;49mexplain_graph(graph\u001b[39m.\u001b[39;49mx, graph\u001b[39m.\u001b[39;49medge_index, edge_attr\u001b[39m=\u001b[39;49mtorch\u001b[39m.\u001b[39;49mzeros(\u001b[39m0\u001b[39;49m))\n",
      "File \u001b[1;32mD:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch_geometric\\nn\\models\\gnn_explainer.py:146\u001b[0m, in \u001b[0;36mGNNExplainer.explain_graph\u001b[1;34m(self, x, edge_index, **kwargs)\u001b[0m\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=142'>143</a>\u001b[0m batch \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mzeros(x\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m], dtype\u001b[39m=\u001b[39m\u001b[39mint\u001b[39m, device\u001b[39m=\u001b[39mx\u001b[39m.\u001b[39mdevice)\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=144'>145</a>\u001b[0m \u001b[39m# Get the initial prediction.\u001b[39;00m\n\u001b[1;32m--> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=145'>146</a>\u001b[0m prediction \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mget_initial_prediction(x, edge_index, batch\u001b[39m=\u001b[39mbatch,\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=146'>147</a>\u001b[0m                                          \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=148'>149</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_initialize_masks(x, edge_index)\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/gnn_explainer.py?line=149'>150</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mto(x\u001b[39m.\u001b[39mdevice)\n",
      "File \u001b[1;32mD:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\autograd\\grad_mode.py:27\u001b[0m, in \u001b[0;36m_DecoratorContextManager.__call__.<locals>.decorate_context\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m     <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/autograd/grad_mode.py?line=23'>24</a>\u001b[0m \u001b[39m@functools\u001b[39m\u001b[39m.\u001b[39mwraps(func)\n\u001b[0;32m     <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/autograd/grad_mode.py?line=24'>25</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdecorate_context\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[0;32m     <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/autograd/grad_mode.py?line=25'>26</a>\u001b[0m     \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mclone():\n\u001b[1;32m---> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/autograd/grad_mode.py?line=26'>27</a>\u001b[0m         \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mD:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch_geometric\\nn\\models\\explainer.py:235\u001b[0m, in \u001b[0;36mExplainer.get_initial_prediction\u001b[1;34m(self, x, edge_index, batch, **kwargs)\u001b[0m\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=230'>231</a>\u001b[0m \u001b[39m@torch\u001b[39m\u001b[39m.\u001b[39mno_grad()\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=231'>232</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mget_initial_prediction\u001b[39m(\u001b[39mself\u001b[39m, x: Tensor, edge_index: Tensor,\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=232'>233</a>\u001b[0m                            batch: Optional[Tensor] \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=233'>234</a>\u001b[0m     \u001b[39mif\u001b[39;00m batch \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m--> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=234'>235</a>\u001b[0m         out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel(x, edge_index, batch, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=235'>236</a>\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[0;32m    <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch_geometric/nn/models/explainer.py?line=236'>237</a>\u001b[0m         out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel(x, edge_index, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mD:\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1105'>1106</a>\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1106'>1107</a>\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1107'>1108</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1108'>1109</a>\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1109'>1110</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39m\u001b[39minput\u001b[39m, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1110'>1111</a>\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   <a href='file:///d%3A/Anaconda3/envs/pytorch/lib/site-packages/torch/nn/modules/module.py?line=1111'>1112</a>\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "\u001b[1;31mTypeError\u001b[0m: forward() got multiple values for argument 'edge_attr'"
     ]
    }
   ],
   "source": [
    "# node_feat_mask, edge_mask = explainer.explain_graph(graph.x, graph.edge_index, batch_index=1)\n",
    "# node_feat_mask, edge_mask = explainer.explain_graph(graph.x, graph.edge_index, batch=3)\n",
    "# node_feat_mask, edge_mask = explainaer.explain_graph(graph.x, graph.edge_index, edge_attr=2)\n",
    "node_feat_mask, edge_mask = explainer.explain_graph(graph.x, graph.edge_index, edge_attr=torch.zeros(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35c5ebe9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([27, 464]) torch.Size([2, 821])\n",
      "torch.Size([27, 492]) torch.Size([2, 1217])\n",
      "torch.Size([27, 485]) torch.Size([2, 925])\n",
      "torch.Size([27, 440]) torch.Size([2, 792])\n",
      "torch.Size([27, 465]) torch.Size([2, 836])\n",
      "torch.Size([27, 468]) torch.Size([2, 855])\n",
      "torch.Size([27, 440]) torch.Size([2, 792])\n",
      "torch.Size([27, 587]) torch.Size([2, 1083])\n",
      "torch.Size([27, 1137]) torch.Size([2, 3479])\n",
      "torch.Size([27, 478]) torch.Size([2, 866])\n",
      "torch.Size([27, 1024]) torch.Size([2, 2837])\n",
      "torch.Size([27, 481]) torch.Size([2, 690])\n",
      "torch.Size([27, 488]) torch.Size([2, 928])\n",
      "torch.Size([27, 322]) torch.Size([2, 354])\n",
      "torch.Size([27, 1122]) torch.Size([2, 2241])\n",
      "torch.Size([27, 559]) torch.Size([2, 982])\n",
      "torch.Size([27, 1643]) torch.Size([2, 4396])\n",
      "torch.Size([27, 568]) torch.Size([2, 1731])\n",
      "torch.Size([27, 463]) torch.Size([2, 816])\n",
      "torch.Size([27, 484]) torch.Size([2, 925])\n"
     ]
    }
   ],
   "source": [
    "# for i in range(len(train_dataset)):\n",
    "for i in range(20):\n",
    "    print(train_dataset[i].x.size(),train_dataset[i].edge_index.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e7d9439",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  0,   0,   0,  ..., 559, 560, 561],\n",
       "        [531, 546, 254,  ..., 120, 522, 523]])"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[i].edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b6e6e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([464, 27])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# train_dataset[i].x\n",
    "# len(train_dataset[i].x)\n",
    "# train_dataset[0].x.size()\n",
    "train_dataset[0].x.size()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17cb122",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[i].edge_attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953eac84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[i].y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46dcc67d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}