02 Graph Attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install of pytorch geometric https://github.com/rusty1s/pytorch_geometric\n",
    "\n",
    "from torch_geometric.datasets import Planetoid\n",
    "from torch_geometric.utils import to_dense_adj, add_self_loops\n",
    "\n",
    "import torch_geometric.transforms as T\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "cora = Planetoid(root='./data', name='Cora')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data(edge_index=[2, 10556], test_mask=[2708], train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])\n",
      "Number of nodes: 2708\n",
      "Number of edges: 10556\n",
      "Number of training nodes: 140\n",
      "Training node label rate: 0.05\n",
      "Is undirected: True\n"
     ]
    }
   ],
   "source": [
    "print(cora.data)\n",
    "print(f'Number of nodes: {cora.data.num_nodes}')\n",
    "print(f'Number of edges: {cora.data.num_edges}')\n",
    "print(f'Number of training nodes: {cora.data.train_mask.sum()}')\n",
    "print(f'Training node label rate: {int(cora.data.train_mask.sum()) / cora.data.num_nodes:.2f}')\n",
    "print(f'Is undirected: {cora.data.is_undirected()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Node features\n",
    "X = cora.data.x\n",
    "\n",
    "# Adjacency matrix + add self-loop\n",
    "A = to_dense_adj(cora.data.edge_index).squeeze()\n",
    "A = A + torch.eye(A.shape[0])\n",
    "\n",
    "# Target\n",
    "y = cora.data.y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 1., 1.,  ..., 0., 0., 0.],\n",
       "        [0., 1., 1.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0.,  ..., 1., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 1., 1.],\n",
       "        [0., 0., 0.,  ..., 0., 1., 1.]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def masked_acc(logits, y, mask):\n",
    "    \n",
    "    a_max = torch.argmax(logits, dim=-1)\n",
    "    \n",
    "    eq = (a_max == y).float() * mask.float()\n",
    "            \n",
    "    return eq.sum() / mask.float().sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def masked_cross_entropy(logits, y, mask):\n",
    "    \n",
    "    masked_y = y.masked_fill(mask == False, value=-1)\n",
    "    \n",
    "    loss = F.cross_entropy(logits, masked_y, ignore_index=-1)\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, A, n_epochs, opt):\n",
    "    \n",
    "    best_acc = 0.\n",
    "\n",
    "    for ep in range(n_epochs):\n",
    "\n",
    "        out = model(X, A)\n",
    "\n",
    "        loss =  masked_cross_entropy(out, y, cora.data.train_mask)\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "\n",
    "        acc = masked_acc(out, y, cora.data.test_mask)\n",
    "\n",
    "        if acc > best_acc:\n",
    "\n",
    "            best_acc = acc\n",
    "\n",
    "            print(f'epoch ={ep: .1f}\\t acc ={best_acc: .1%}')\n",
    "            \n",
    "    return best_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAT(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_feat=1433, hid_feat=64, num_classes=7):\n",
    "        \n",
    "        super().__init__()\n",
    "                        \n",
    "        self.W_att_1 = nn.Linear(in_feat, 3 * hid_feat, bias=False)\n",
    "        \n",
    "        self.W_att_2 = nn.Linear(hid_feat, 3 * num_classes, bias=False)\n",
    "        \n",
    "    def compute_attention(self, H, A, W):\n",
    "        \"\"\"\n",
    "        Dot-product attention\n",
    "        \n",
    "        H: Hidden features [N, D]\n",
    "        A: Adjacency matrix [N, N]\n",
    "        W: layer for attention [D, 3 * D]\n",
    "        \"\"\"\n",
    "                \n",
    "        H = W(H) # [N, 3 * dim]\n",
    "                                \n",
    "        query, key, value = torch.chunk(H, 3, dim=-1)\n",
    "        \n",
    "        d_model = query.size(-1)\n",
    "    \n",
    "        score = query @ key.transpose(-1, -2) / np.sqrt(d_model) # [N, N]\n",
    "                \n",
    "        score = score.masked_fill(A == 0, float('-inf')) # [N, N]\n",
    "                        \n",
    "        normalized_score = F.softmax(score, dim=1) # [N, N]\n",
    "                        \n",
    "        return normalized_score @ value # [N, D]\n",
    "        \n",
    "        \n",
    "    def forward(self, X, A):\n",
    "        \"\"\"\n",
    "        inputs:\n",
    "            X: Node features of shape [N, in_feat]\n",
    "            A: Adjacency matrix of shape [N, N]\n",
    "        returns:\n",
    "            O: Logits of shape [N, num_classes]\n",
    "        \"\"\"\n",
    "        \n",
    "        H = self.compute_attention(X, A, self.W_att_1)\n",
    "        \n",
    "        H = torch.relu(H)\n",
    "        \n",
    "        O = self.compute_attention(H, A, self.W_att_2)\n",
    "        \n",
    "        return O"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GAT(hid_feat=128)\n",
    "\n",
    "opt = torch.optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch = 0.0\t acc = 27.5%\n",
      "epoch = 1.0\t acc = 47.6%\n",
      "epoch = 2.0\t acc = 62.3%\n",
      "epoch = 3.0\t acc = 70.6%\n",
      "epoch = 4.0\t acc = 74.3%\n",
      "epoch = 5.0\t acc = 75.2%\n",
      "epoch = 6.0\t acc = 76.3%\n",
      "epoch = 7.0\t acc = 77.6%\n",
      "epoch = 8.0\t acc = 78.5%\n",
      "epoch = 9.0\t acc = 79.1%\n",
      "epoch = 10.0\t acc = 79.2%\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-16-75a5c6184c4e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-8-a3edae643a4d>\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(model, A, n_epochs, opt)\u001b[0m\n\u001b[1;32m     11\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m         \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0macc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmasked_acc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcora\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest_mask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/optim/adam.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m    117\u001b[0m                    \u001b[0mgroup\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'lr'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    118\u001b[0m                    \u001b[0mgroup\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'weight_decay'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 119\u001b[0;31m                    \u001b[0mgroup\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'eps'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    120\u001b[0m                    )\n\u001b[1;32m    121\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/optim/functional.py\u001b[0m in \u001b[0;36madam\u001b[0;34m(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, amsgrad, beta1, beta2, lr, weight_decay, eps)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     85\u001b[0m         \u001b[0;31m# Decay the first and second moment running average coefficient\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 86\u001b[0;31m         \u001b[0mexp_avg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mbeta1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     87\u001b[0m         \u001b[0mexp_avg_sq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maddcmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mbeta2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mamsgrad\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train_model(model, A, 100, opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GAT with PyG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GATConv, GATv2Conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATpyG(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_feat=1433, hid_feat=64, num_classes=7):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.gat_1 = GATv2Conv(in_channels=in_feat, out_channels=hid_feat)\n",
    "        \n",
    "        self.gat_2 = GATv2Conv(in_channels=hid_feat, out_channels=num_classes)\n",
    "        \n",
    "    def forward(self, data):\n",
    "        \n",
    "        X, edge_index = data.x, data.edge_index\n",
    "        # edge_index of shape [2, num_edges]\n",
    "        # X of shape [N, in_feat]\n",
    "        \n",
    "        \n",
    "        # the gcn layer takes the node features and the edge_index tensor\n",
    "        H = self.gat_1(X, edge_index) # [N, hid_feat]\n",
    "        \n",
    "        H = torch.relu(H)\n",
    "        \n",
    "        O = self.gat_2(H, edge_index) # [N, num_classes]\n",
    "        \n",
    "        return O"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model_pyg(model, n_epochs, opt):\n",
    "    \n",
    "    best_acc = 0.\n",
    "\n",
    "    for ep in range(n_epochs):\n",
    "\n",
    "        out = model(cora.data)\n",
    "\n",
    "        loss =  masked_cross_entropy(out, y, cora.data.train_mask)\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "\n",
    "        acc = masked_acc(out, y, cora.data.test_mask)\n",
    "\n",
    "        if acc > best_acc:\n",
    "\n",
    "            best_acc = acc\n",
    "\n",
    "            print(f'epoch ={ep: .1f}\\t acc ={best_acc: .1%}')\n",
    "            \n",
    "    return best_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg_gcn = GATpyG()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_2 = torch.optim.Adam(pyg_gcn.parameters(), lr=0.002)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch = 0.0\t acc = 7.8%\n",
      "epoch = 1.0\t acc = 23.1%\n",
      "epoch = 2.0\t acc = 39.4%\n",
      "epoch = 3.0\t acc = 52.4%\n",
      "epoch = 4.0\t acc = 62.0%\n",
      "epoch = 5.0\t acc = 67.4%\n",
      "epoch = 6.0\t acc = 69.8%\n",
      "epoch = 7.0\t acc = 71.6%\n",
      "epoch = 8.0\t acc = 73.6%\n",
      "epoch = 9.0\t acc = 74.1%\n",
      "epoch = 10.0\t acc = 74.7%\n",
      "epoch = 11.0\t acc = 75.1%\n",
      "epoch = 12.0\t acc = 75.8%\n",
      "epoch = 13.0\t acc = 76.3%\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor(0.7630)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_model_pyg(pyg_gcn, 100, opt_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py37_UZA",
   "language": "python",
   "name": "py37_uza"
  },
  "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.7.9"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}