minimal_graph_attention_network.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CmaJDQfwjpV1",
        "outputId": "842ab54e-f736-4d5f-eb19-efc008c4632a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Output shape: torch.Size([34, 2])\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import networkx as nx\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "class GraphAttentionLayer(nn.Module):\n",
        "    def __init__(self, input_dim, output_dim, dropout_rate=0.5, alpha=0.2):\n",
        "        super(GraphAttentionLayer, self).__init__()\n",
        "        self.dropout = nn.Dropout(dropout_rate)\n",
        "        self.alpha = alpha\n",
        "        self.weights = nn.Parameter(torch.FloatTensor(input_dim, output_dim))\n",
        "        self.bias = nn.Parameter(torch.FloatTensor(output_dim))\n",
        "\n",
        "        # Attention mechanisms\n",
        "        self.attention = nn.Parameter(torch.FloatTensor(2 * output_dim, 1))\n",
        "\n",
        "        # Initialize parameters\n",
        "        nn.init.xavier_uniform_(self.weights.data)\n",
        "        nn.init.xavier_uniform_(self.attention.data)\n",
        "        nn.init.constant_(self.bias.data, 0.0)\n",
        "\n",
        "    def forward(self, x, adjacency):\n",
        "        x = torch.matmul(x, self.weights)\n",
        "        x = torch.matmul(adjacency, x)\n",
        "\n",
        "        # Attention mechanism\n",
        "        num_nodes = x.size()[0]\n",
        "        attention_input = torch.cat([x.repeat(1, num_nodes).view(num_nodes * num_nodes, -1),\n",
        "                                     x.repeat(num_nodes, 1)], dim=1).view(num_nodes, -1, 2 * x.size(1))\n",
        "\n",
        "        attention_weights = F.leaky_relu(torch.matmul(attention_input, self.attention), negative_slope=self.alpha)\n",
        "        attention_weights = F.softmax(attention_weights, dim=1)\n",
        "        attention_weights = self.dropout(attention_weights)\n",
        "\n",
        "        x = torch.matmul(attention_weights.transpose(1, 2), x)\n",
        "        x = x.squeeze()\n",
        "        x = x + self.bias\n",
        "        return x\n",
        "\n",
        "\n",
        "class GraphAttentionNetwork(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim, num_heads, dropout_rate=0.5, alpha=0.2):\n",
        "        super(GraphAttentionNetwork, self).__init__()\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.num_heads = num_heads\n",
        "\n",
        "        self.attentions = nn.ModuleList()\n",
        "        for _ in range(num_heads):\n",
        "            self.attentions.append(GraphAttentionLayer(input_dim, hidden_dim, dropout_rate, alpha))\n",
        "\n",
        "        self.out_att = GraphAttentionLayer(hidden_dim * num_heads, output_dim, dropout_rate, alpha)\n",
        "\n",
        "    def forward(self, x, adjacency):\n",
        "        x = torch.cat([att(x, adjacency) for att in self.attentions], dim=1)\n",
        "        x = self.out_att(x, adjacency)\n",
        "        return x\n",
        "\n",
        "\n",
        "# Create a sample graph using NetworkX\n",
        "graph = nx.karate_club_graph()\n",
        "\n",
        "# Generate adjacency matrix\n",
        "adjacency = nx.adjacency_matrix(graph)\n",
        "adjacency = torch.tensor(adjacency.todense(), dtype=torch.float32)\n",
        "\n",
        "# Generate node features\n",
        "features = np.eye(graph.number_of_nodes(), dtype=np.float32)\n",
        "features = torch.tensor(features, dtype=torch.float32)\n",
        "\n",
        "# Define model parameters\n",
        "input_dim = features.shape[1]\n",
        "hidden_dim = 8\n",
        "output_dim = 2\n",
        "num_heads = 2\n",
        "\n",
        "# Create GAT model\n",
        "model = GraphAttentionNetwork(input_dim, hidden_dim, output_dim, num_heads)\n",
        "\n",
        "# Perform forward pass\n",
        "output = model(features, adjacency)\n",
        "print(\"Output shape:\", output.shape)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "EhJ4cJ1fjpyZ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}