EvolveGCN with KD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "EvolveGCN(KD).ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HgkABpYeWtZr"
      },
      "source": [
        "## **Pytorch Geometric Environment Setting**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6E7ouY2uW3iK"
      },
      "source": [
        "!python -m pip install --upgrade pip\n",
        "!python -m pip install pip==20.2.4"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vXTY0Zp6W6IP"
      },
      "source": [
        "import torch\n",
        "\n",
        "def format_pytorch_version(version):\n",
        "  return version.split('+')[0]\n",
        "\n",
        "TORCH_version = torch.__version__\n",
        "TORCH = format_pytorch_version(TORCH_version)\n",
        "\n",
        "def format_cuda_version(version):\n",
        "  return 'cu' + version.replace('.', '')\n",
        "\n",
        "CUDA_version = torch.version.cuda\n",
        "CUDA = format_cuda_version(CUDA_version)\n",
        "\n",
        "!pip install torch-scatter     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-sparse      -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-cluster     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-geometric "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ROfB6OG2XSgJ"
      },
      "source": [
        "## **Library import**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q2rj8T7aXU9d"
      },
      "source": [
        "import numpy as np\n",
        "import networkx as nx\n",
        "import os\n",
        "import pandas as pd\n",
        "from sklearn.metrics import f1_score, precision_score, recall_score\n",
        "from sklearn.manifold import TSNE\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.nn import Embedding\n",
        "from torch.nn import Parameter\n",
        "from torch_geometric.data import Data,DataLoader\n",
        "from torch_geometric.nn import GCNConv\n",
        "from torch_geometric.utils.convert import to_networkx\n",
        "from torch_geometric.utils import to_undirected"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouxS_THKXcVX"
      },
      "source": [
        "## **Dataset preparation**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yrZKWYKWXfjn"
      },
      "source": [
        "os.environ['KAGGLE_USERNAME'] = \"karthikapv\" # username from the json file\n",
        "os.environ['KAGGLE_KEY'] = \"cc11b8fcbb2e177d31cd566bbabe382a\" # key from the json file\n",
        "!kaggle datasets download -d ellipticco/elliptic-data-set\n",
        "!unzip elliptic-data-set.zip\n",
        "!mkdir elliptic_bitcoin_dataset_cont"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bPRoEr2EXgNa"
      },
      "source": [
        "import time\n",
        "import tracemalloc\n",
        "\n",
        "\n",
        "def get_memory_and_execution_time_details(func, is_teacher):\n",
        "    tracemalloc.start()\n",
        "    start_time = time.time()\n",
        "    func(teacher=is_teacher)\n",
        "    exec_time = time.time() - start_time\n",
        "    print(\"Model Evaluation Time: \")\n",
        "    print(exec_time)\n",
        "    current, peak = tracemalloc.get_traced_memory()\n",
        "    print(f\"Current memory usage is {current / 10 ** 3}KB; Peak was {peak / 10 ** 3}KB\")\n",
        "    tracemalloc.stop()\n",
        "\n",
        "    return current, peak, exec_time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "369ogvPoX0An"
      },
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "from torch_geometric.data import Data, DataLoader\n",
        "from torch_geometric.utils import to_undirected\n",
        "\n",
        "\n",
        "def train_test_split():\n",
        "    df_edge = pd.read_csv('elliptic_bitcoin_dataset/elliptic_txs_edgelist.csv')\n",
        "    df_class = pd.read_csv('elliptic_bitcoin_dataset/elliptic_txs_classes.csv')\n",
        "    df_features = pd.read_csv('elliptic_bitcoin_dataset/elliptic_txs_features.csv', header=None)\n",
        "\n",
        "    # Setting Column name\n",
        "    df_features.columns = ['id', 'time step'] + [f'trans_feat_{i}' for i in range(93)] + [f'agg_feat_{i}' for i in\n",
        "                                                                                          range(72)]\n",
        "\n",
        "    print('Number of edges: {}'.format(len(df_edge)))\n",
        "    df_edge.head()\n",
        "\n",
        "    # Get Node Index\n",
        "\n",
        "    all_nodes = list(\n",
        "        set(df_edge['txId1']).union(set(df_edge['txId2'])).union(set(df_class['txId'])).union(set(df_features['id'])))\n",
        "    nodes_df = pd.DataFrame(all_nodes, columns=['id']).reset_index()\n",
        "\n",
        "    print('Number of nodes: {}'.format(len(nodes_df)))\n",
        "    nodes_df.head()\n",
        "\n",
        "    # Fix id index\n",
        "\n",
        "    df_edge = df_edge.join(nodes_df.rename(columns={'id': 'txId1'}).set_index('txId1'), on='txId1', how='inner') \\\n",
        "        .join(nodes_df.rename(columns={'id': 'txId2'}).set_index('txId2'), on='txId2', how='inner', rsuffix='2') \\\n",
        "        .drop(columns=['txId1', 'txId2']) \\\n",
        "        .rename(columns={'index': 'txId1', 'index2': 'txId2'})\n",
        "    df_edge.head()\n",
        "\n",
        "    df_class = df_class.join(nodes_df.rename(columns={'id': 'txId'}).set_index('txId'), on='txId', how='inner') \\\n",
        "        .drop(columns=['txId']).rename(columns={'index': 'txId'})[['txId', 'class']]\n",
        "    df_class.head()\n",
        "\n",
        "    df_features = df_features.join(nodes_df.set_index('id'), on='id', how='inner') \\\n",
        "        .drop(columns=['id']).rename(columns={'index': 'id'})\n",
        "    df_features = df_features[['id'] + list(df_features.drop(columns=['id']).columns)]\n",
        "    df_features.head()\n",
        "\n",
        "    df_edge_time = df_edge.join(df_features[['id', 'time step']].rename(columns={'id': 'txId1'}).set_index('txId1'),\n",
        "                                on='txId1', how='left', rsuffix='1') \\\n",
        "        .join(df_features[['id', 'time step']].rename(columns={'id': 'txId2'}).set_index('txId2'), on='txId2', how='left',\n",
        "              rsuffix='2')\n",
        "    df_edge_time['is_time_same'] = df_edge_time['time step'] == df_edge_time['time step2']\n",
        "    df_edge_time_fin = df_edge_time[['txId1', 'txId2', 'time step']].rename(\n",
        "        columns={'txId1': 'source', 'txId2': 'target', 'time step': 'time'})\n",
        "\n",
        "    # Create csv from Dataframe\n",
        "\n",
        "    df_features.drop(columns=['time step']).to_csv('elliptic_bitcoin_dataset_cont/elliptic_txs_features.csv', index=False, header=None)\n",
        "    df_class.rename(columns={'txId': 'nid', 'class': 'label'})[['nid', 'label']].sort_values(by='nid').to_csv(\n",
        "        'elliptic_bitcoin_dataset_cont/elliptic_txs_classes.csv', index=False, header=None)\n",
        "    df_features[['id', 'time step']].rename(columns={'id': 'nid', 'time step': 'time'})[['nid', 'time']].sort_values(\n",
        "        by='nid').to_csv('elliptic_bitcoin_dataset_cont/elliptic_txs_nodetime.csv', index=False, header=None)\n",
        "    df_edge_time_fin[['source', 'target', 'time']].to_csv('elliptic_bitcoin_dataset_cont/elliptic_txs_edgelist_timed.csv', index=False,\n",
        "                                                          header=None)\n",
        "\n",
        "    # Graph Preprocessing\n",
        "\n",
        "    node_label = df_class.rename(columns={'txId': 'nid', 'class': 'label'})[['nid', 'label']].sort_values(by='nid').merge(\n",
        "        df_features[['id', 'time step']].rename(columns={'id': 'nid', 'time step': 'time'}), on='nid', how='left')\n",
        "    node_label['label'] = node_label['label'].apply(lambda x: '3' if x == 'unknown' else x).astype(int) - 1\n",
        "    node_label.head()\n",
        "\n",
        "    merged_nodes_df = node_label.merge(\n",
        "        df_features.rename(columns={'id': 'nid', 'time step': 'time'}).drop(columns=['time']), on='nid', how='left')\n",
        "    merged_nodes_df.head()\n",
        "\n",
        "    train_dataset = []\n",
        "    test_dataset = []\n",
        "\n",
        "    num_node_features = 0\n",
        "    for i in range(49):\n",
        "        nodes_df_tmp = merged_nodes_df[merged_nodes_df['time'] == i + 1].reset_index()\n",
        "        nodes_df_tmp['index'] = nodes_df_tmp.index\n",
        "        df_edge_tmp = df_edge_time_fin.join(\n",
        "            nodes_df_tmp.rename(columns={'nid': 'source'})[['source', 'index']].set_index('source'), on='source',\n",
        "            how='inner') \\\n",
        "            .join(nodes_df_tmp.rename(columns={'nid': 'target'})[['target', 'index']].set_index('target'), on='target',\n",
        "                  how='inner', rsuffix='2') \\\n",
        "            .drop(columns=['source', 'target']) \\\n",
        "            .rename(columns={'index': 'source', 'index2': 'target'})\n",
        "        x = torch.tensor(np.array(nodes_df_tmp.sort_values(by='index').drop(columns=['index', 'nid', 'label'])),\n",
        "                         dtype=torch.float)\n",
        "        edge_index = torch.tensor(np.array(df_edge_tmp[['source', 'target']]).T, dtype=torch.long)\n",
        "        edge_index = to_undirected(edge_index)\n",
        "        mask = nodes_df_tmp['label'] != 2\n",
        "        y = torch.tensor(np.array(nodes_df_tmp['label']), dtype=torch.long)\n",
        "\n",
        "        data = Data(x=x, edge_index=edge_index, mask=mask, y=y)\n",
        "        num_node_features = data.num_node_features\n",
        "        if i + 1 < 35:\n",
        "            train_dataset.append(data)\n",
        "        else:\n",
        "            test_dataset.append(data)\n",
        "\n",
        "    train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
        "    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "    return train_loader, test_loader, num_node_features"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7J-_8gx-YVkt"
      },
      "source": [
        "## **Model**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fcqq_N1DYW0q"
      },
      "source": [
        "import torch\n",
        "from torch.nn import Parameter\n",
        "from torch_geometric.nn import GCNConv\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class GCN(torch.nn.Module):\n",
        "    def __init__(self, num_node_features, hidden_channels, use_skip=False):\n",
        "        super(GCN, self).__init__()\n",
        "        self.conv1 = GCNConv(num_node_features, hidden_channels[0])\n",
        "        self.conv2 = GCNConv(hidden_channels[0], 2)\n",
        "        self.use_skip = use_skip\n",
        "        if self.use_skip:\n",
        "            self.weight = torch.nn.init.xavier_normal_(Parameter(torch.Tensor(num_node_features, 2)))\n",
        "\n",
        "    def forward(self, data):\n",
        "        x = self.conv1(data.x, data.edge_index)\n",
        "        x = x.relu()\n",
        "        x = F.dropout(x, p=0.5, training=self.training)\n",
        "        x = self.conv2(x, data.edge_index)\n",
        "        if self.use_skip:\n",
        "            x = F.softmax(x + torch.matmul(x, self.weight), dim=-1)\n",
        "        else:\n",
        "            x = F.softmax(x, dim=-1)\n",
        "        return x\n",
        "\n",
        "    def embed(self, data):\n",
        "        x = self.conv1(data.x, data.edge_index)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4HeVoalPYaX_"
      },
      "source": [
        "import torch\n",
        "from torch.nn import LSTM\n",
        "from torch_geometric.nn import GCNConv\n",
        "\n",
        "\n",
        "class EvolveGCNO(torch.nn.Module):\n",
        "    r\"\"\"An implementation of the Evolving Graph Convolutional without Hidden Layer.\n",
        "    For details see this paper: `\"EvolveGCN: Evolving Graph Convolutional \n",
        "    Networks for Dynamic Graph.\" <https://arxiv.org/abs/1902.10191>`_\n",
        "\n",
        "    Args:\n",
        "        in_channels (int): Number of filters.\n",
        "        improved (bool, optional): If set to :obj:`True`, the layer computes\n",
        "            :math:`\\mathbf{\\hat{A}}` as :math:`\\mathbf{A} + 2\\mathbf{I}`.\n",
        "            (default: :obj:`False`)\n",
        "        cached (bool, optional): If set to :obj:`True`, the layer will cache\n",
        "            the computation of :math:`\\mathbf{\\hat{D}}^{-1/2} \\mathbf{\\hat{A}}\n",
        "            \\mathbf{\\hat{D}}^{-1/2}` on first execution, and will use the\n",
        "            cached version for further executions.\n",
        "            This parameter should only be set to :obj:`True` in transductive\n",
        "            learning scenarios. (default: :obj:`False`)\n",
        "        normalize (bool, optional): Whether to add self-loops and apply\n",
        "            symmetric normalization. (default: :obj:`True`)\n",
        "        add_self_loops (bool, optional): If set to :obj:`False`, will not add\n",
        "            self-loops to the input graph. (default: :obj:`True`)\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels: int, improved: bool=False, cached: bool=False,\n",
        "                 normalize: bool=True, add_self_loops: bool=True):\n",
        "        super(EvolveGCNO, self).__init__()\n",
        "\n",
        "        self.in_channels = in_channels\n",
        "        self.improved = improved\n",
        "        self.cached = cached\n",
        "        self.normalize = normalize\n",
        "        self.add_self_loops = add_self_loops\n",
        "        self._create_layers()\n",
        "\n",
        "\n",
        "    def _create_layers(self):\n",
        "\n",
        "        self.recurrent_layer = LSTM(input_size = self.in_channels,\n",
        "                                    hidden_size = self.in_channels,\n",
        "                                    num_layers = 1)\n",
        "\n",
        "\n",
        "        self.conv_layer = GCNConv(in_channels = self.in_channels,\n",
        "                                  out_channels = self.in_channels,\n",
        "                                  improved = self.improved,\n",
        "                                  cached = self.cached,\n",
        "                                  normalize = self.normalize,\n",
        "                                  add_self_loops = self.add_self_loops,\n",
        "                                  bias = False)\n",
        "\n",
        "    def forward(self, X: torch.FloatTensor, edge_index: torch.LongTensor, \n",
        "                edge_weight: torch.FloatTensor=None) -> torch.FloatTensor:\n",
        "        \"\"\"\n",
        "        Making a forward pass.\n",
        "\n",
        "        Arg types:\n",
        "            * **X** *(PyTorch Float Tensor)* - Node embedding.\n",
        "            * **edge_index** *(PyTorch Long Tensor)* - Graph edge indices.\n",
        "            * **edge_weight** *(PyTorch Float Tensor, optional)* - Edge weight vector.\n",
        "\n",
        "        Return types:\n",
        "            * **X** *(PyTorch Float Tensor)* - Output matrix for all nodes.\n",
        "        \"\"\"\n",
        "        W = self.conv_layer.weight[None, :, :]\n",
        "        W, _ = self.recurrent_layer(W)\n",
        "        self.conv_layer.weight = torch.nn.Parameter(W.squeeze())\n",
        "        X = self.conv_layer(X, edge_index, edge_weight)\n",
        "        return X"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WYkUn5pqYhLJ"
      },
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch.nn import Parameter\n",
        "#from torch_geometric_temporal.nn.recurrent import EvolveGCNO\n",
        "\n",
        "\n",
        "class RecurrentGCN(torch.nn.Module):\n",
        "    def __init__(self, node_features, num_classes, dropout_rate=0.5):\n",
        "        super(RecurrentGCN, self).__init__()\n",
        "        self.dropout_rate = dropout_rate\n",
        "        self.recurrent_1 = EvolveGCNO(node_features, num_classes)\n",
        "        self.recurrent_2 = EvolveGCNO(node_features, num_classes)\n",
        "        self.linear = torch.nn.Linear(node_features, num_classes)\n",
        "\n",
        "    def forward(self, data):\n",
        "        x = self.recurrent_1(data.x, data.edge_index)\n",
        "        x = F.relu(x)\n",
        "        x = F.dropout(x, p=self.dropout_rate, training=self.training)\n",
        "        x = self.recurrent_2(x, data.edge_index)\n",
        "        x = F.relu(x)\n",
        "        x = F.dropout(x, p=self.dropout_rate, training=self.training)\n",
        "        x = self.linear(x)\n",
        "        return F.log_softmax(x, dim=-1)\n",
        "\n",
        "    def embed(self, data):\n",
        "        x = self.recurrent_1(data.x, data.edge_index)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N0BZrjOXYlhm"
      },
      "source": [
        "import os\n",
        "from copy import deepcopy\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from sklearn.metrics import f1_score, precision_score, recall_score\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "\n",
        "\n",
        "class BaseClass:\n",
        "    \"\"\"\n",
        "    Basic implementation of a general Knowledge Distillation framework\n",
        "    :param teacher_model (torch.nn.Module): Teacher model\n",
        "    :param student_model (torch.nn.Module): Student model\n",
        "    :param train_loader (torch.utils.data.DataLoader): Dataloader for training\n",
        "    :param val_loader (torch.utils.data.DataLoader): Dataloader for validation/testing\n",
        "    :param optimizer_teacher (torch.optim.*): Optimizer used for training teacher\n",
        "    :param optimizer_student (torch.optim.*): Optimizer used for training student\n",
        "    :param loss_fn (torch.nn.Module): Loss Function used for distillation\n",
        "    :param temp (float): Temperature parameter for distillation\n",
        "    :param distil_weight (float): Weight paramter for distillation loss\n",
        "    :param device (str): Device used for training; 'cpu' for cpu and 'cuda' for gpu\n",
        "    :param log (bool): True if logging required\n",
        "    :param logdir (str): Directory for storing logs\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "            self,\n",
        "            teacher_model,\n",
        "            student_model,\n",
        "            train_loader,\n",
        "            val_loader,\n",
        "            optimizer_teacher,\n",
        "            optimizer_student,\n",
        "            loss_fn=nn.KLDivLoss(),\n",
        "            temp=20.0,\n",
        "            distil_weight=0.5,\n",
        "            device=\"cpu\",\n",
        "            log=False,\n",
        "            logdir=\"./Experiments\",\n",
        "    ):\n",
        "\n",
        "        self.train_loader = train_loader\n",
        "        self.val_loader = val_loader\n",
        "        self.optimizer_teacher = optimizer_teacher\n",
        "        self.optimizer_student = optimizer_student\n",
        "        self.temp = temp\n",
        "        self.distil_weight = distil_weight\n",
        "        self.log = log\n",
        "        self.logdir = logdir\n",
        "\n",
        "        if self.log:\n",
        "            self.writer = SummaryWriter(logdir)\n",
        "\n",
        "        try:\n",
        "            torch.Tensor(0).to(device)\n",
        "            self.device = device\n",
        "        except:\n",
        "            print(\n",
        "                \"Either an invalid device or CUDA is not available. Defaulting to CPU.\"\n",
        "            )\n",
        "            self.device = torch.device(\"cpu\")\n",
        "\n",
        "        try:\n",
        "            self.teacher_model = teacher_model.to(self.device)\n",
        "        except:\n",
        "            print(\"Warning!!! Teacher is NONE.\")\n",
        "        self.student_model = student_model.to(self.device)\n",
        "        try:\n",
        "            self.loss_fn = loss_fn.to(self.device)\n",
        "            self.ce_fn = nn.CrossEntropyLoss().to(self.device)\n",
        "        except:\n",
        "            self.loss_fn = loss_fn\n",
        "            self.ce_fn = nn.CrossEntropyLoss()\n",
        "            print(\"Warning: Loss Function can't be moved to device.\")\n",
        "\n",
        "    def train_teacher(\n",
        "            self,\n",
        "            epochs=10,\n",
        "            plot_losses=True,\n",
        "            save_model=True,\n",
        "            save_model_pth=\"./models/teacher.pt\",\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Function that will be training the teacher\n",
        "        :param epochs (int): Number of epochs you want to train the teacher\n",
        "        :param plot_losses (bool): True if you want to plot the losses\n",
        "        :param save_model (bool): True if you want to save the teacher model\n",
        "        :param save_model_pth (str): Path where you want to store the teacher model\n",
        "        \"\"\"\n",
        "        self.teacher_model.train()\n",
        "        loss_arr = []\n",
        "        illicit_f1_arr = []\n",
        "        micro_avg_f1_arr = []\n",
        "        illicit_precision_arr = []\n",
        "        micro_avg_precision_arr = []\n",
        "        illicit_recall_arr = []\n",
        "        micro_avg_recall_arr = []\n",
        "        length_of_dataset = len(self.train_loader.dataset)\n",
        "        best_acc = 0.0\n",
        "        self.best_teacher_model_weights = deepcopy(self.teacher_model.state_dict())\n",
        "\n",
        "        save_dir = os.path.dirname(save_model_pth)\n",
        "        if not os.path.exists(save_dir):\n",
        "            os.makedirs(save_dir)\n",
        "\n",
        "        print(\"Training Teacher... \")\n",
        "\n",
        "        for ep in range(epochs):\n",
        "            epoch_loss = 0.0\n",
        "            correct = 0\n",
        "            torch.manual_seed(ep)\n",
        "            np.random.seed(42)\n",
        "            torch.backends.cudnn.deterministic = True\n",
        "            torch.backends.cudnn.benchmark = False\n",
        "            for data in self.train_loader:\n",
        "\n",
        "                data.x = data.x.to(self.device)\n",
        "                label = data.y.to(self.device)\n",
        "                mask = data.mask\n",
        "\n",
        "                out = self.teacher_model(data)\n",
        "\n",
        "                if isinstance(out, tuple):\n",
        "                    out = out[0]\n",
        "\n",
        "                pred = out.argmax(dim=1, keepdim=True)\n",
        "                correct += pred.eq(label.view_as(pred)).sum().item()\n",
        "                illicit_f1_arr.append(f1_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_f1_arr.append(f1_score(pred[mask], label[mask], average='micro'))\n",
        "                illicit_precision_arr.append(precision_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_precision_arr.append(precision_score(pred[mask], label[mask], average='micro'))\n",
        "                illicit_recall_arr.append(recall_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_recall_arr.append(recall_score(pred[mask], label[mask], average='micro'))\n",
        "\n",
        "                loss = self.ce_fn(out[mask], label[mask])\n",
        "\n",
        "                self.optimizer_teacher.zero_grad()\n",
        "                loss.backward()\n",
        "                self.optimizer_teacher.step()\n",
        "\n",
        "                epoch_loss += loss\n",
        "\n",
        "            epoch_acc = correct / length_of_dataset\n",
        "            if epoch_acc > best_acc:\n",
        "                best_acc = epoch_acc\n",
        "                self.best_teacher_model_weights = deepcopy(\n",
        "                    self.teacher_model.state_dict()\n",
        "                )\n",
        "\n",
        "            if self.log:\n",
        "                self.writer.add_scalar(\"Training loss/Teacher\", epoch_loss, epochs)\n",
        "                self.writer.add_scalar(\"Training accuracy/Teacher\", epoch_acc, epochs)\n",
        "\n",
        "            loss_arr.append(epoch_loss)\n",
        "            print(\n",
        "                'Epoch: {:1d}, Epoch Loss: {:.4f}, Illicit Precision: {:.4f}, Illicit Recall: '\n",
        "                '{:.4f}, Illicit f1: {:.4f}, F1: {:.4f}, Precision: {:.4f}, Recall: {:.4f}' \\\n",
        "                    .format(ep + 1, epoch_loss, np.mean(illicit_precision_arr),\n",
        "                            np.mean(illicit_recall_arr), np.mean(illicit_f1_arr), np.mean(micro_avg_f1_arr),\n",
        "                            np.mean(micro_avg_precision_arr), np.mean(micro_avg_recall_arr)))\n",
        "\n",
        "            self.post_epoch_call(ep)\n",
        "\n",
        "        self.teacher_model.load_state_dict(self.best_teacher_model_weights)\n",
        "        if save_model:\n",
        "            torch.save(self.teacher_model.state_dict(), save_model_pth)\n",
        "        if plot_losses:\n",
        "            plt.plot(loss_arr)\n",
        "\n",
        "    def _train_student(\n",
        "            self,\n",
        "            epochs=10,\n",
        "            plot_losses=True,\n",
        "            save_model=True,\n",
        "            save_model_pth=\"./models/student.pt\",\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Function to train student model - for internal use only.\n",
        "        :param epochs (int): Number of epochs you want to train the teacher\n",
        "        :param plot_losses (bool): True if you want to plot the losses\n",
        "        :param save_model (bool): True if you want to save the student model\n",
        "        :param save_model_pth (str): Path where you want to save the student model\n",
        "        \"\"\"\n",
        "        self.teacher_model.eval()\n",
        "        self.student_model.train()\n",
        "        loss_arr = []\n",
        "        illicit_f1_arr = []\n",
        "        micro_avg_f1_arr = []\n",
        "        illicit_precision_arr = []\n",
        "        micro_avg_precision_arr = []\n",
        "        illicit_recall_arr = []\n",
        "        micro_avg_recall_arr = []\n",
        "        length_of_dataset = len(self.train_loader.dataset)\n",
        "        best_acc = 0.0\n",
        "        self.best_student_model_weights = deepcopy(self.student_model.state_dict())\n",
        "\n",
        "        save_dir = os.path.dirname(save_model_pth)\n",
        "        if not os.path.exists(save_dir):\n",
        "            os.makedirs(save_dir)\n",
        "\n",
        "        print(\"Training Student...\")\n",
        "\n",
        "        for ep in range(epochs):\n",
        "            epoch_loss = 0.0\n",
        "            correct = 0\n",
        "            torch.manual_seed(ep)\n",
        "            np.random.seed(ep)\n",
        "            torch.backends.cudnn.deterministic = True\n",
        "            torch.backends.cudnn.benchmark = False\n",
        "            for data in self.train_loader:\n",
        "\n",
        "                data.x = data.x.to(self.device)\n",
        "                label = data.y.to(self.device)\n",
        "                mask = data.mask\n",
        "\n",
        "                student_out = self.student_model(data)\n",
        "                teacher_out = self.teacher_model(data)\n",
        "\n",
        "                loss = self.calculate_kd_loss(student_out[mask], teacher_out[mask], label[mask])\n",
        "\n",
        "                if isinstance(student_out, tuple):\n",
        "                    student_out = student_out[0]\n",
        "\n",
        "                pred = student_out.argmax(dim=1, keepdim=True)\n",
        "                correct += pred.eq(label.view_as(pred)).sum().item()\n",
        "                illicit_f1_arr.append(f1_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_f1_arr.append(f1_score(pred[mask], label[mask], average='micro'))\n",
        "                illicit_precision_arr.append(precision_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_precision_arr.append(precision_score(pred[mask], label[mask], average='micro'))\n",
        "                illicit_recall_arr.append(recall_score(pred[mask], label[mask], pos_label=1))\n",
        "                micro_avg_recall_arr.append(recall_score(pred[mask], label[mask], average='micro'))\n",
        "\n",
        "                self.optimizer_student.zero_grad()\n",
        "                loss.backward()\n",
        "                self.optimizer_student.step()\n",
        "\n",
        "                epoch_loss += loss.item()\n",
        "\n",
        "            epoch_acc = correct / length_of_dataset\n",
        "            if epoch_acc > best_acc:\n",
        "                best_acc = epoch_acc\n",
        "                self.best_student_model_weights = deepcopy(\n",
        "                    self.student_model.state_dict()\n",
        "                )\n",
        "\n",
        "            if self.log:\n",
        "                self.writer.add_scalar(\"Training loss/Student\", epoch_loss, epochs)\n",
        "                self.writer.add_scalar(\"Training accuracy/Student\", epoch_acc, epochs)\n",
        "\n",
        "            loss_arr.append(epoch_loss)\n",
        "            print(\n",
        "                'Epoch: {:1d}, Epoch Loss: {:.4f}, Illicit Precision: {:.4f}, Illicit Recall: '\n",
        "                '{:.4f}, Illicit f1: {:.4f}, F1: {:.4f}, Precision: {:.4f}, Recall: {:.4f}' \\\n",
        "                    .format(ep + 1, epoch_loss, np.mean(illicit_precision_arr),\n",
        "                            np.mean(illicit_recall_arr), np.mean(illicit_f1_arr), np.mean(micro_avg_f1_arr),\n",
        "                            np.mean(micro_avg_precision_arr), np.mean(micro_avg_recall_arr)))\n",
        "\n",
        "        self.student_model.load_state_dict(self.best_student_model_weights)\n",
        "        if save_model:\n",
        "            torch.save(self.student_model.state_dict(), save_model_pth)\n",
        "        if plot_losses:\n",
        "            plt.plot(loss_arr)\n",
        "\n",
        "    def train_student(\n",
        "            self,\n",
        "            epochs=10,\n",
        "            plot_losses=True,\n",
        "            save_model=True,\n",
        "            save_model_pth=\"./models/student.pt\",\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Function that will be training the student\n",
        "        :param epochs (int): Number of epochs you want to train the teacher\n",
        "        :param plot_losses (bool): True if you want to plot the losses\n",
        "        :param save_model (bool): True if you want to save the student model\n",
        "        :param save_model_pth (str): Path where you want to save the student model\n",
        "        \"\"\"\n",
        "        self._train_student(epochs, plot_losses, save_model, save_model_pth)\n",
        "\n",
        "    def calculate_kd_loss(self, y_pred_student, y_pred_teacher, y_true):\n",
        "        \"\"\"\n",
        "        Custom loss function to calculate the KD loss for various implementations\n",
        "        :param y_pred_student (Tensor): Predicted outputs from the student network\n",
        "        :param y_pred_teacher (Tensor): Predicted outputs from the teacher network\n",
        "        :param y_true (Tensor): True labels\n",
        "        \"\"\"\n",
        "\n",
        "        raise NotImplementedError\n",
        "\n",
        "    def _evaluate_model(self, model, verbose=False):\n",
        "        \"\"\"\n",
        "        Evaluate the given model's accuaracy over val set.\n",
        "        For internal use only.\n",
        "        :param model (nn.Module): Model to be used for evaluation\n",
        "        :param verbose (bool): Display Accuracy\n",
        "        \"\"\"\n",
        "        model.eval()\n",
        "        length_of_dataset = len(self.val_loader.dataset)\n",
        "        correct = 0\n",
        "        outputs = []\n",
        "        illicit_f1_arr = []\n",
        "        micro_avg_f1_arr = []\n",
        "        illicit_precision_arr = []\n",
        "        micro_avg_precision_arr = []\n",
        "        illicit_recall_arr = []\n",
        "        micro_avg_recall_arr = []\n",
        "\n",
        "        seed_val = 35\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for data in self.train_loader:\n",
        "\n",
        "                torch.manual_seed(seed_val)\n",
        "                np.random.seed(seed_val)\n",
        "                torch.backends.cudnn.deterministic = True\n",
        "                torch.backends.cudnn.benchmark = False\n",
        "\n",
        "                data.x = data.x.to(self.device)\n",
        "                target = data.y.to(self.device)\n",
        "                mask = data.mask\n",
        "\n",
        "                output = model(data)\n",
        "\n",
        "                if isinstance(output, tuple):\n",
        "                    output = output[0]\n",
        "                outputs.append(output)\n",
        "\n",
        "                pred = output.argmax(dim=1, keepdim=True)\n",
        "                correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "                accuracy = correct / length_of_dataset\n",
        "                illicit_f1_arr.append(f1_score(pred[mask], target[mask], pos_label=1))\n",
        "                micro_avg_f1_arr.append(f1_score(pred[mask], target[mask], average='micro'))\n",
        "                illicit_precision_arr.append(precision_score(pred[mask], target[mask], pos_label=1))\n",
        "                micro_avg_precision_arr.append(precision_score(pred[mask], target[mask], average='micro'))\n",
        "                illicit_recall_arr.append(recall_score(pred[mask], target[mask], pos_label=1))\n",
        "                micro_avg_recall_arr.append(recall_score(pred[mask], target[mask], average='micro'))\n",
        "\n",
        "                if verbose:\n",
        "                    print(\"-\" * 80)\n",
        "                    print(f\"Iteration: {seed_val-34}\")\n",
        "                    print(\"-\" * 80)\n",
        "                    print(\"Illicit F1: {:.4f}\".format(f1_score(pred[mask], target[mask], pos_label=1)))\n",
        "                    print(\"Illicit Precision: {:.4f}\".format(precision_score(pred[mask], target[mask], pos_label=1)))\n",
        "                    print(\"Illicit Recall: {:.4f}\".format(recall_score(pred[mask], target[mask], pos_label=1)))\n",
        "                    print(\"Micro Avg F1: {:.4f}\".format(f1_score(pred[mask], target[mask], average='micro')))\n",
        "                    print(\"Micro Avg Precision: {:.4f}\".format(precision_score(pred[mask], target[mask], average='micro')))\n",
        "                    print(\"Micro Avg Recall: {:.4f}\".format(recall_score(pred[mask], target[mask], average='micro')))\n",
        "\n",
        "                seed_val += 1\n",
        "\n",
        "        print(\"-\" * 80)\n",
        "        print(\"-\" * 80)\n",
        "        print(\"Final Result\")\n",
        "        print(\"-\" * 80)\n",
        "        print(\"-\" * 80)\n",
        "        \n",
        "        print(\"Illicit F1: {:.4f}\".format(np.mean(illicit_f1_arr)))\n",
        "        print(\"Illicit Precision: {:.4f}\".format(np.mean(illicit_precision_arr)))\n",
        "        print(\"Illicit Recall: {:.4f}\".format(np.mean(illicit_recall_arr)))\n",
        "        print(\"Micro Avg F1: {:.4f}\".format(np.mean(micro_avg_f1_arr)))\n",
        "        print(\"Micro Avg Precision: {:.4f}\".format(np.mean(micro_avg_precision_arr)))\n",
        "        print(\"Micro Avg Recall: {:.4f}\".format(np.mean(micro_avg_recall_arr)))\n",
        "        return outputs, accuracy\n",
        "\n",
        "    def evaluate(self, teacher=False):\n",
        "        \"\"\"\n",
        "        Evaluate method for printing accuracies of the trained network\n",
        "        :param teacher (bool): True if you want accuracy of the teacher network\n",
        "        \"\"\"\n",
        "        if teacher:\n",
        "            model = deepcopy(self.teacher_model).to(self.device)\n",
        "        else:\n",
        "            model = deepcopy(self.student_model).to(self.device)\n",
        "        _, accuracy = self._evaluate_model(model=model, verbose=False)\n",
        "\n",
        "        return accuracy\n",
        "\n",
        "    def get_parameters(self):\n",
        "        \"\"\"\n",
        "        Get the number of parameters for the teacher and the student network\n",
        "        \"\"\"\n",
        "        teacher_params = sum(p.numel() for p in self.teacher_model.parameters())\n",
        "        student_params = sum(p.numel() for p in self.student_model.parameters())\n",
        "\n",
        "        print(\"-\" * 80)\n",
        "        print(f\"Total parameters for the teacher network are: {teacher_params}\")\n",
        "        print(f\"Total parameters for the student network are: {student_params}\")\n",
        "\n",
        "    def post_epoch_call(self, epoch):\n",
        "        \"\"\"\n",
        "        Any changes to be made after an epoch is completed.\n",
        "        :param epoch (int) : current epoch number\n",
        "        :return            : nothing (void)\n",
        "        \"\"\"\n",
        "\n",
        "        pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xtEDu4-kYyaa"
      },
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "\n",
        "class VanillaKD(BaseClass):\n",
        "    \"\"\"\n",
        "    Original implementation of Knowledge distillation from the paper \"Distilling the\n",
        "    Knowledge in a Neural Network\" https://arxiv.org/pdf/1503.02531.pdf\n",
        "    :param teacher_model (torch.nn.Module): Teacher model\n",
        "    :param student_model (torch.nn.Module): Student model\n",
        "    :param train_loader (torch.utils.data.DataLoader): Dataloader for training\n",
        "    :param val_loader (torch.utils.data.DataLoader): Dataloader for validation/testing\n",
        "    :param optimizer_teacher (torch.optim.*): Optimizer used for training teacher\n",
        "    :param optimizer_student (torch.optim.*): Optimizer used for training student\n",
        "    :param loss_fn (torch.nn.Module):  Calculates loss during distillation\n",
        "    :param temp (float): Temperature parameter for distillation\n",
        "    :param distil_weight (float): Weight paramter for distillation loss\n",
        "    :param device (str): Device used for training; 'cpu' for cpu and 'cuda' for gpu\n",
        "    :param log (bool): True if logging required\n",
        "    :param logdir (str): Directory for storing logs\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        teacher_model,\n",
        "        student_model,\n",
        "        train_loader,\n",
        "        val_loader,\n",
        "        optimizer_teacher,\n",
        "        optimizer_student,\n",
        "        loss_fn=nn.MSELoss(),\n",
        "        temp=20.0,\n",
        "        distil_weight=0.5,\n",
        "        device=\"cpu\",\n",
        "        log=False,\n",
        "        logdir=\"./Experiments\",\n",
        "    ):\n",
        "        super(VanillaKD, self).__init__(\n",
        "            teacher_model,\n",
        "            student_model,\n",
        "            train_loader,\n",
        "            val_loader,\n",
        "            optimizer_teacher,\n",
        "            optimizer_student,\n",
        "            loss_fn,\n",
        "            temp,\n",
        "            distil_weight,\n",
        "            device,\n",
        "            log,\n",
        "            logdir,\n",
        "        )\n",
        "\n",
        "    def calculate_kd_loss(self, y_pred_student, y_pred_teacher, y_true):\n",
        "        \"\"\"\n",
        "        Function used for calculating the KD loss during distillation\n",
        "        :param y_pred_student (torch.FloatTensor): Prediction made by the student model\n",
        "        :param y_pred_teacher (torch.FloatTensor): Prediction made by the teacher model\n",
        "        :param y_true (torch.FloatTensor): Original label\n",
        "        \"\"\"\n",
        "\n",
        "        soft_teacher_out = F.softmax(y_pred_teacher / self.temp, dim=1)\n",
        "        soft_student_out = F.softmax(y_pred_student / self.temp, dim=1)\n",
        "\n",
        "        loss = (1 - self.distil_weight) * F.cross_entropy(y_pred_student, y_true)\n",
        "        loss += (self.distil_weight * self.temp * self.temp) * self.loss_fn(\n",
        "            soft_teacher_out, soft_student_out\n",
        "        )\n",
        "\n",
        "        return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r0AVpxV9Y7j_"
      },
      "source": [
        "## **Train**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FRMP4RBxY7FR"
      },
      "source": [
        "import time\n",
        "\n",
        "import torch.optim as optim\n",
        "from torch_geometric.data import Data, DataLoader\n",
        "\n",
        "\n",
        "train_loader, test_loader, num_node_features = train_test_split()\n",
        "\n",
        "student_model = GCN(num_node_features=num_node_features, hidden_channels=[100])\n",
        "\n",
        "lr = 0.001\n",
        "weight_decay = 0.005\n",
        "\n",
        "student_optimizer = optim.Adam(student_model.parameters(), lr=lr, weight_decay=weight_decay, amsgrad=True)\n",
        "\n",
        "# Distilled EvolveGCN\n",
        "\n",
        "epochs = 10\n",
        "\n",
        "teacher_model_evolvegcn = RecurrentGCN(node_features=num_node_features, num_classes=2)\n",
        "\n",
        "teacher_optimizer_evolvegcn = optim.Adam(teacher_model_evolvegcn.parameters(), lr=lr, weight_decay=weight_decay,\n",
        "                                         amsgrad=True)\n",
        "\n",
        "distiller_evolvegcn = VanillaKD(teacher_model_evolvegcn, student_model, train_loader, test_loader,\n",
        "                                   teacher_optimizer_evolvegcn, student_optimizer)\n",
        "distiller_evolvegcn.train_teacher(epochs=epochs, plot_losses=True, save_model=True,\n",
        "                                  save_model_pth='./models/evolvegcn_teacher.pt')  # Train the teacher network\n",
        "\n",
        "distiller_evolvegcn.train_student(epochs=epochs, plot_losses=True, save_model=True,\n",
        "                                  save_model_pth='./models/evolvegcn_student.pt')  # Train the student network"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "klSi1WDiZO1-"
      },
      "source": [
        "## **Results**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CSlyTm9AZSIt"
      },
      "source": [
        "import time\n",
        "\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "\n",
        "\n",
        "train_loader, test_loader, num_node_features = train_test_split()\n",
        "\n",
        "student_model = GCN(num_node_features=num_node_features, hidden_channels=[100])\n",
        "\n",
        "lr = 10e-5\n",
        "weight_decay = 5e-4\n",
        "\n",
        "student_optimizer = optim.Adam(student_model.parameters(), lr=lr, weight_decay=weight_decay, amsgrad=True)\n",
        "\n",
        "# Distilled EvolveGCN\n",
        "\n",
        "epochs = 10\n",
        "\n",
        "teacher_model_evolvegcn = RecurrentGCN(node_features=num_node_features, num_classes=2)\n",
        "teacher_model_evolvegcn.load_state_dict(torch.load(\"./models/evolvegcn_teacher.pt\"))\n",
        "\n",
        "teacher_optimizer_evolvegcn = optim.Adam(teacher_model_evolvegcn.parameters(), lr=lr, weight_decay=weight_decay,\n",
        "                                         amsgrad=True)\n",
        "\n",
        "student_model.load_state_dict(torch.load(\"./models/evolvegcn_student.pt\"))\n",
        "\n",
        "distiller_evolvegcn = VanillaKD(teacher_model_evolvegcn, student_model, train_loader, test_loader,\n",
        "                                   teacher_optimizer_evolvegcn, student_optimizer)\n",
        "\n",
        "get_memory_and_execution_time_details(distiller_evolvegcn.evaluate, True)  # Evaluate the teacher network\n",
        "\n",
        "get_memory_and_execution_time_details(distiller_evolvegcn.evaluate, False)  # Evaluate the student network\n",
        "\n",
        "distiller_evolvegcn.get_parameters()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}