Q2_GNN/gat_old.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "resources:\n",
    "- https://www.youtube.com/watch?v=SnRfBfXwLuY\n",
    "- https://arxiv.org/pdf/1710.10903.pdf\n",
    "- https://www.youtube.com/watch?v=CwsPoa7z2c8\n",
    "- https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.GATConv.html\n",
    "- https://github.com/pyg-team/pytorch_geometric/blob/master/examples/ppi.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.datasets import PPI\n",
    "from torch_geometric.data import DataLoader\n",
    "import torch_geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATLayer(nn.Module):\n",
    "    def __init__(self, g, in_dim, out_dim):\n",
    "        super(GATLayer, self).__init__()\n",
    "        self.g = g\n",
    "        # equation (1)\n",
    "        self.fc = nn.Linear(in_dim, out_dim, bias=False)\n",
    "        # equation (2)\n",
    "        self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        \"\"\"Reinitialize learnable parameters.\"\"\"\n",
    "        gain = nn.init.calculate_gain('relu')\n",
    "        nn.init.xavier_normal_(self.fc.weight, gain=gain)\n",
    "        nn.init.xavier_normal_(self.attn_fc.weight, gain=gain)\n",
    "\n",
    "    def edge_attention(self, edges):\n",
    "        # edge UDF for equation (2)\n",
    "        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)\n",
    "        a = self.attn_fc(z2)\n",
    "        return {'e': F.leaky_relu(a)}\n",
    "\n",
    "    def message_func(self, edges):\n",
    "        # message UDF for equation (3) & (4)\n",
    "        return {'z': edges.src['z'], 'e': edges.data['e']}\n",
    "\n",
    "    def reduce_func(self, nodes):\n",
    "        # reduce UDF for equation (3) & (4)\n",
    "        # equation (3)\n",
    "        alpha = F.softmax(nodes.mailbox['e'], dim=1)\n",
    "        # equation (4)\n",
    "        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)\n",
    "        return {'h': h}\n",
    "\n",
    "    def forward(self, h):\n",
    "        # equation (1)\n",
    "        z = self.fc(h)\n",
    "        self.g.ndata['z'] = z\n",
    "        # equation (2)\n",
    "        self.g.apply_edges(self.edge_attention)\n",
    "        # equation (3) & (4)\n",
    "        self.g.update_all(self.message_func, self.reduce_func)\n",
    "        return self.g.ndata.pop('h')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadGATLayer(nn.Module):\n",
    "    def __init__(self, g, in_dim, out_dim, num_heads, merge='cat'):\n",
    "        super(MultiHeadGATLayer, self).__init__()\n",
    "        self.heads = nn.ModuleList()\n",
    "        for i in range(num_heads):\n",
    "            self.heads.append(GATLayer(g, in_dim, out_dim))\n",
    "        self.merge = merge\n",
    "\n",
    "    def forward(self, h):\n",
    "        head_outs = [attn_head(h) for attn_head in self.heads]\n",
    "        if self.merge == 'cat':\n",
    "            # concat on the output feature dimension (dim=1)\n",
    "            return torch.cat(head_outs, dim=1)\n",
    "        else:\n",
    "            # merge using average\n",
    "            return torch.mean(torch.stack(head_outs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAT(nn.Module):\n",
    "    def __init__(self, g, in_dim, hidden_dim, out_dim, num_heads):\n",
    "        super(GAT, self).__init__()\n",
    "        self.layer1 = MultiHeadGATLayer(g, in_dim, hidden_dim, num_heads)\n",
    "        # Be aware that the input dimension is hidden_dim*num_heads since\n",
    "        # multiple head outputs are concatenated together. Also, only\n",
    "        # one attention head in the output layer.\n",
    "        self.layer2 = MultiHeadGATLayer(g, hidden_dim * num_heads, out_dim, 1)\n",
    "\n",
    "    def forward(self, h):\n",
    "        h = self.layer1(h)\n",
    "        h = F.elu(h)\n",
    "        h = self.layer2(h)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Classes in Cora: 7\n",
      "Number of Node Features in Cora: 1433\n"
     ]
    }
   ],
   "source": [
    "from torch_geometric.data import Data\n",
    "from torch_geometric.nn import GATConv\n",
    "from torch_geometric.datasets import Planetoid\n",
    "import torch_geometric.transforms as T\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "name_data = 'Cora'\n",
    "dataset = Planetoid(root= '/tmp/' + name_data, name = name_data)\n",
    "dataset.transform = T.NormalizeFeatures()\n",
    "\n",
    "print(f\"Number of Classes in {name_data}:\", dataset.num_classes)\n",
    "print(f\"Number of Node Features in {name_data}:\", dataset.num_node_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train(train_dataloader, val_dataloader, device, model):\n",
    "    # define loss function and optimizer\n",
    "    loss_fcn = nn.BCEWithLogitsLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=0)\n",
    "\n",
    "    # training loop\n",
    "    for epoch in range(400):\n",
    "        model.train()\n",
    "        logits = []\n",
    "        total_loss = 0\n",
    "        # mini-batch loop\n",
    "        for batch_id, batched_graph in enumerate(train_dataloader):\n",
    "            batched_graph = batched_graph.to(device)\n",
    "            features = batched_graph.x.float()\n",
    "            labels = batched_graph.y.float()\n",
    "            logits = model(batched_graph, features)\n",
    "            loss = loss_fcn(logits, labels)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += loss.item()\n",
    "        print(\n",
    "            \"Epoch {:05d} | Loss {:.4f} |\".format(\n",
    "                epoch, total_loss / (batch_id + 1)\n",
    "            )\n",
    "        )\n",
    "\n",
    "        if (epoch + 1) % 5 == 0:\n",
    "            avg_score = evaluate_in_batches(\n",
    "                val_dataloader, device, model\n",
    "            )  # evaluate F1-score instead of loss\n",
    "            print(\n",
    "                \"                            Acc. (F1-score) {:.4f} \".format(\n",
    "                    avg_score\n",
    "                )\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def evaluate(g, features, labels, model):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        output = model(g, features)\n",
    "        pred = np.where(output.data.cpu().numpy() >= 0, 1, 0)\n",
    "        score = f1_score(labels.data.cpu().numpy(), pred, average=\"micro\")\n",
    "        return score\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_in_batches(dataloader, device, model):\n",
    "    total_score = 0\n",
    "    for batch_id, batched_graph in enumerate(dataloader):\n",
    "        batched_graph = batched_graph.to(device)\n",
    "        features = batched_graph.ndata[\"feat\"]\n",
    "        labels = batched_graph.ndata[\"label\"]\n",
    "        score = evaluate(batched_graph, features, labels, model)\n",
    "        total_score += score\n",
    "    return total_score / (batch_id + 1)  # return average score\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from torch_geometric.datasets import PPI\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "path = 'data/'\n",
    "train_dataset = PPI(path, split='train')\n",
    "val_dataset = PPI(path, split='val')\n",
    "test_dataset = PPI(path, split='test')\n",
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "121\n",
      "Training...\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "forward() takes 2 positional arguments but 3 were given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[111], line 16\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;66;03m# model training\u001b[39;00m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTraining...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 16\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# test the model\u001b[39;00m\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTesting...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[106], line 16\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(train_dataloader, val_dataloader, device, model)\u001b[0m\n\u001b[1;32m     14\u001b[0m features \u001b[38;5;241m=\u001b[39m batched_graph\u001b[38;5;241m.\u001b[39mx\u001b[38;5;241m.\u001b[39mfloat()\n\u001b[1;32m     15\u001b[0m labels \u001b[38;5;241m=\u001b[39m batched_graph\u001b[38;5;241m.\u001b[39my\u001b[38;5;241m.\u001b[39mfloat()\n\u001b[0;32m---> 16\u001b[0m logits \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatched_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     17\u001b[0m loss \u001b[38;5;241m=\u001b[39m loss_fcn(logits, labels)\n\u001b[1;32m     18\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "\u001b[0;31mTypeError\u001b[0m: forward() takes 2 positional arguments but 3 were given"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "print()\n",
    "features = train_dataset.num_classes\n",
    " \n",
    "g = torch_geometric.utils.to_networkx(train_dataset[0])\n",
    "\n",
    "# create GAT model\n",
    "in_size = features\n",
    "print(in_size)\n",
    "out_size = train_dataset.num_classes\n",
    "model = GAT(g=g, in_dim=in_size, hidden_dim=256, out_dim=out_size, num_heads=2).to(device)\n",
    "\n",
    "# model training\n",
    "print(\"Training...\")\n",
    "train(train_loader, val_loader, device, model)\n",
    "\n",
    "# test the model\n",
    "print(\"Testing...\")\n",
    "avg_score = evaluate_in_batches(test_loader, device, model)\n",
    "print(\"Test Accuracy (F1-score) {:.4f}\".format(avg_score))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the PPI dataset\n",
    "dataset = PPI(root='data/')  # Replace with the desired path\n",
    "data = dataset[0]  # PPI consists of a single graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAT(torch.nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, num_heads):\n",
    "        super(GAT, self).__init__()\n",
    "        self.conv1 = MultiHeadGATLayer(in_channels, hidden_channels, num_heads=num_heads, out_dim=hidden_channels)\n",
    "        self.conv2 = MultiHeadGATLayer(hidden_channels * num_heads, out_channels, num_heads=1, out_dim=out_channels)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = F.elu(self.conv1(x, edge_index))\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        return F.log_softmax(x, dim=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the model instance\n",
    "model = GAT(dataset.num_features, 8, dataset.num_classes, num_heads=2)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "forward() takes 2 positional arguments but 3 were given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[69], line 8\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m epoch \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m3\u001b[39m:\n\u001b[1;32m      6\u001b[0m     t0 \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m----> 8\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      9\u001b[0m loss \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mnll_loss(out[data\u001b[38;5;241m.\u001b[39mtrain_mask], data\u001b[38;5;241m.\u001b[39my[data\u001b[38;5;241m.\u001b[39mtrain_mask])\n\u001b[1;32m     11\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Cell \u001b[0;32mIn[66], line 9\u001b[0m, in \u001b[0;36mGAT.forward\u001b[0;34m(self, x, edge_index)\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, edge_index):\n\u001b[1;32m      8\u001b[0m     x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mdropout(x, p\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.6\u001b[39m, training\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining)\n\u001b[0;32m----> 9\u001b[0m     x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39melu(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv1\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     10\u001b[0m     x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mdropout(x, p\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.6\u001b[39m, training\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining)\n\u001b[1;32m     11\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv2(x, edge_index)\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "\u001b[0;31mTypeError\u001b[0m: forward() takes 2 positional arguments but 3 were given"
     ]
    }
   ],
   "source": [
    "import time\n",
    "# Main training loop\n",
    "dur = []\n",
    "for epoch in range(30):\n",
    "    if epoch >= 3:\n",
    "        t0 = time.time()\n",
    "\n",
    "    out = model(data.x, data.edge_index)\n",
    "    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if epoch >= 3:\n",
    "        dur.append(time.time() - t0)\n",
    "\n",
    "    print(\"Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f}\".format(\n",
    "        epoch, loss.item(), np.mean(dur)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.utils import add_self_loops, degree\n",
    "\n",
    "class GATConv(MessagePassing):\n",
    "    def __init__(self, in_channels, out_channels, heads=1, concat=True):\n",
    "        super(GATConv, self).__init__(aggr='add')\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.heads = heads\n",
    "        self.concat = concat\n",
    "        self.weight = torch.nn.Parameter(torch.Tensor(in_channels, out_channels * heads))\n",
    "        self.att = torch.nn.Parameter(torch.Tensor(1, heads, 2 * out_channels))\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        torch.nn.init.xavier_uniform_(self.weight)\n",
    "        torch.nn.init.xavier_uniform_(self.att)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))\n",
    "        x = torch.matmul(x, self.weight)\n",
    "        x = x.view(-1, self.heads, self.out_channels)\n",
    "        return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x)\n",
    "\n",
    "    def message(self, x_i, x_j, edge_index_i, size):\n",
    "        alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)\n",
    "        alpha = F.leaky_relu(alpha, negative_slope=0.2)\n",
    "        alpha = softmax(alpha, edge_index_i, size)\n",
    "        return x_j * alpha.view(-1, self.heads, 1)\n",
    "\n",
    "    def update(self, aggr_out):\n",
    "        if self.concat:\n",
    "            aggr_out = aggr_out.view(-1, self.out_channels * self.heads)\n",
    "        else:\n",
    "            aggr_out = aggr_out.mean(dim=1)\n",
    "        return aggr_out\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.utils import softmax\n",
    "\n",
    "class GATConv2(torch.nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, heads=1, concat=True):\n",
    "        super(GATConv, self).__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.heads = heads\n",
    "        self.concat = concat\n",
    "        self.weight = torch.nn.Parameter(torch.Tensor(in_channels, out_channels * heads))\n",
    "        self.att = torch.nn.Parameter(torch.Tensor(1, heads, 2 * out_channels))\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        torch.nn.init.xavier_uniform_(self.weight)\n",
    "        torch.nn.init.xavier_uniform_(self.att)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        print(\"Input tensor shape:\", x.shape)\n",
    "        print(\"Edge index shape:\", edge_index.shape)\n",
    "        edge_index, _ = torch.geometric.utils.add_self_loops(edge_index, num_nodes=x.size(0))\n",
    "        x = torch.matmul(x, self.weight)\n",
    "        x = x.view(-1, self.heads, self.out_channels)\n",
    "        row, col = edge_index\n",
    "        alpha = (torch.cat([x[row], x[col]], dim=-1) * self.att).sum(dim=-1)\n",
    "        alpha = F.leaky_relu(alpha, negative_slope=0.2)\n",
    "        alpha = softmax(alpha, row, num_nodes=x.size(0))\n",
    "        aggr_out = torch.matmul(alpha.unsqueeze(1), x)\n",
    "        if self.concat:\n",
    "            aggr_out = aggr_out.view(-1, self.out_channels * self.heads)\n",
    "        else:\n",
    "            aggr_out = aggr_out.mean(dim=1)\n",
    "        return aggr_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Parameter\n",
    "from torch_geometric.nn.inits import glorot, zeros\n",
    "\n",
    "class GATConv(torch.nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, heads=1, negative_slope=0.2, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.heads = heads\n",
    "        self.negative_slope = negative_slope\n",
    "        self.dropout = dropout\n",
    "\n",
    "        self.lin = torch.nn.Linear(in_channels, heads * out_channels, bias=False)\n",
    "        self.att_src = Parameter(torch.empty(1, heads, out_channels))\n",
    "        self.att_dst = Parameter(torch.empty(1, heads, out_channels))\n",
    "\n",
    "        self.bias = Parameter(torch.empty(heads * out_channels))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        glorot(self.lin.weight)\n",
    "        glorot(self.att_src)\n",
    "        glorot(self.att_dst)\n",
    "        zeros(self.bias)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = self.lin(x).view(-1, self.heads, self.out_channels)\n",
    "\n",
    "        alpha_src = (x * self.att_src).sum(dim=-1)\n",
    "        alpha_dst = (x * self.att_dst).sum(dim=-1)\n",
    "\n",
    "        alpha = alpha_src[edge_index[0]] + alpha_dst[edge_index[1]]        \n",
    "        alpha = F.leaky_relu(alpha, self.negative_slope)\n",
    "        alpha = softmax(alpha, edge_index[0], num_nodes=x.size(0))\n",
    "        alpha = F.dropout(alpha, p=self.dropout, training=self.training)\n",
    "\n",
    "        out = self.propagate(alpha, x, edge_index)\n",
    "        out = out.view(-1, self.heads * self.out_channels)\n",
    "        out = out + self.bias\n",
    "\n",
    "        return out\n",
    "\n",
    "    def propagate(self, alpha, x, edge_index):\n",
    "        num_nodes = x.size(0)\n",
    "        return alpha.view(-1, self.heads, 1) * x[edge_index[1]].reshape(-1, self.heads, self.out_channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class GraphAttentionLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    Simple GAT layer, similar to https://arxiv.org/abs/1710.10903\n",
    "    \"\"\"\n",
    "    def __init__(self, in_features, out_features, dropout, alpha, concat=True):\n",
    "        super(GraphAttentionLayer, self).__init__()\n",
    "        self.dropout = dropout\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.alpha = alpha\n",
    "        self.concat = concat\n",
    "\n",
    "        self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))\n",
    "        nn.init.xavier_uniform_(self.W.data, gain=1.414)\n",
    "        self.a = nn.Parameter(torch.empty(size=(2*out_features, 1)))\n",
    "        nn.init.xavier_uniform_(self.a.data, gain=1.414)\n",
    "\n",
    "        self.leakyrelu = nn.LeakyReLU(self.alpha)\n",
    "\n",
    "    def forward(self, h, adj):\n",
    "        Wh = torch.mm(h, self.W) # h.shape: (N, in_features), Wh.shape: (N, out_features)\n",
    "        e = self._prepare_attentional_mechanism_input(Wh)\n",
    "\n",
    "        zero_vec = -9e15*torch.ones_like(e)\n",
    "        attention = torch.where(adj > 0, e, zero_vec)\n",
    "        attention = F.softmax(attention, dim=1)\n",
    "        attention = F.dropout(attention, self.dropout, training=self.training)\n",
    "        h_prime = torch.matmul(attention, Wh)\n",
    "\n",
    "        if self.concat:\n",
    "            return F.elu(h_prime)\n",
    "        else:\n",
    "            return h_prime\n",
    "\n",
    "    def _prepare_attentional_mechanism_input(self, Wh):\n",
    "        # Wh.shape (N, out_feature)\n",
    "        # self.a.shape (2 * out_feature, 1)\n",
    "        # Wh1&2.shape (N, 1)\n",
    "        # e.shape (N, N)\n",
    "        Wh1 = torch.matmul(Wh, self.a[:self.out_features, :])\n",
    "        Wh2 = torch.matmul(Wh, self.a[self.out_features:, :])\n",
    "        # broadcast add\n",
    "        e = Wh1 + Wh2.T\n",
    "        return self.leakyrelu(e)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')'\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class GAT(nn.Module):\n",
    "    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):\n",
    "        \"\"\"Dense version of GAT.\"\"\"\n",
    "        super(GAT, self).__init__()\n",
    "        self.dropout = dropout\n",
    "\n",
    "        self.attentions = [GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in range(nheads)]\n",
    "        for i, attention in enumerate(self.attentions):\n",
    "            self.add_module('attention_{}'.format(i), attention)\n",
    "\n",
    "        self.out_att = GraphAttentionLayer(nhid * nheads, nclass, dropout=dropout, alpha=alpha, concat=False)\n",
    "\n",
    "    def forward(self, x, adj):\n",
    "        x = F.dropout(x, self.dropout, training=self.training)\n",
    "        x = torch.cat([att(x, adj) for att in self.attentions], dim=1)\n",
    "        x = F.dropout(x, self.dropout, training=self.training)\n",
    "        x = F.elu(self.out_att(x, adj))\n",
    "        return F.log_softmax(x, dim=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.datasets import PPI\n",
    "from torch_geometric.data import DataLoader\n",
    "# from torch_geometric.nn import GATConv\n",
    "from sklearn import metrics\n",
    "\n",
    "path = osp.join(osp.dirname(osp.realpath('data/')), 'data', 'PPI')\n",
    "train_dataset = PPI(path, split='train')\n",
    "val_dataset = PPI(path, split='test')\n",
    "test_dataset = PPI(path, split='test')\n",
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)\n",
    "\n",
    "## get number of features, labels, etc.\n",
    "data = train_dataset[0]\n",
    "num_features = train_dataset.num_features\n",
    "num_classes = train_dataset.num_classes\n",
    "num_nodes = data.num_nodes\n",
    "num_edges = data.num_edges\n",
    "\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = GATConv(train_dataset.num_features, 256, heads=4)\n",
    "        self.lin1 = torch.nn.Linear(train_dataset.num_features, 4 * 256)\n",
    "        self.conv2 = GATConv(4 * 256, 256, heads=4)\n",
    "        self.lin2 = torch.nn.Linear(4 * 256, 4 * 256)\n",
    "        self.conv3 = GATConv(\n",
    "            4 * 256, train_dataset.num_classes, heads=6, concat=False)\n",
    "        self.lin3 = torch.nn.Linear(4 * 256, train_dataset.num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = F.elu(self.conv1(x, edge_index) + self.lin1(x))\n",
    "        x = F.elu(self.conv2(x, edge_index) + self.lin2(x))\n",
    "        x = self.conv3(x, edge_index) + self.lin3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = Net().to(device)\n",
    "loss_op = torch.nn.BCEWithLogitsLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.005)\n",
    "\n",
    "model = GAT(nfeat=train_dataset.num_features, \n",
    "                nhid=8, \n",
    "                nclass=int(labels.max()) + 1, \n",
    "                dropout=0.6, \n",
    "                nheads=8, \n",
    "                alpha=0.2)\n",
    "\n",
    "\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "\n",
    "    total_loss = 0\n",
    "    for data in train_loader:\n",
    "        num_graphs = data.num_graphs\n",
    "        data.batch = None\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        loss = loss_op(model(data.x, data.edge_index), data.y)\n",
    "        total_loss += loss.item() * num_graphs\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return total_loss / len(train_loader.dataset)\n",
    "\n",
    "\n",
    "def test(loader):\n",
    "    model.eval()\n",
    "\n",
    "    total_micro_f1 = 0\n",
    "    for data in loader:\n",
    "        with torch.no_grad():\n",
    "            out = model(data.x.to(device), data.edge_index.to(device))\n",
    "        pred = (out > 0).float().cpu()\n",
    "        micro_f1 = metrics.f1_score(data.y, pred, average='micro')\n",
    "        total_micro_f1 += micro_f1 * data.num_graphs\n",
    "    return total_micro_f1 / len(loader.dataset)\n",
    "\n",
    "\n",
    "for epoch in range(1, 101):\n",
    "    loss = train()\n",
    "    acc = test(val_loader)\n",
    "    print('Epoch: {:02d}, Loss: {:.4f}, Acc: {:.4f}'.format(epoch, loss, acc))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}