path: root/transformer_shortest_paths.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "collapsed_sections": [
        "LPphBnKR-aWF",
        "gKt-yIpDebF1"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Step 0: Imports"
      ],
      "metadata": {
        "id": "LPphBnKR-aWF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "import numpy as np\n",
        "from collections import deque\n",
        "import pickle\n",
        "from tqdm import tqdm\n",
        "np.random.seed(42)\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import pickle\n",
        "from math import sqrt\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "import matplotlib.pyplot as plt\n",
        "torch.manual_seed(42)\n",
        "\n",
        "print(\"imports complete\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ge5QvElvhCOw",
        "outputId": "8d2f46b5-22c3-42a7-ecef-ce014d7ec2c9"
      },
      "execution_count": 97,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "imports complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 1: Generate synthetic data"
      ],
      "metadata": {
        "id": "gKt-yIpDebF1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1IbzGIWseK3E",
        "outputId": "835a0467-e1d3-414b-99cb-b8a1b368dd86"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 55/55 [00:00<00:00, 100.02it/s]\n"
          ]
        }
      ],
      "source": [
        "PAD_TOKEN = 0\n",
        "MAX_VTXS = 63\n",
        "# vertices are labelled 1,2,...,63\n",
        "# we also have a padding token which is 0.\n",
        "\n",
        "INF = MAX_VTXS # represents unreachability\n",
        "SEQ_LEN = 128\n",
        "\n",
        "# original task data\n",
        "NTRAIN1 = 10000\n",
        "# the data will be edge lists\n",
        "# like this: [1 3 1 5 2 4 0 0 0 0]\n",
        "# this represents edges (1,3), (1,5) (2,4)\n",
        "# (the zeros are just padding tokens)\n",
        "\n",
        "# the label is the shortest distance from vtx 1 to vtx 2\n",
        "# or \"INF\" if no path exists\n",
        "\n",
        "# fine tuning data\n",
        "NTRAIN2 = 2000\n",
        "# I haven't totally figured out how to do the fine tuning yet.\n",
        "# So don't worry about this yet.\n",
        "\n",
        "def random_graph(n):\n",
        "    assert n >= 4\n",
        "    edge_list = []\n",
        "    adjacencies = [set() for _ in range(n+1)]\n",
        "\n",
        "    indices = np.random.randint(n, size=(2*n))+1\n",
        "    for i in range(0, len(indices), 2):\n",
        "        u = indices[i]\n",
        "        v = indices[i + 1]\n",
        "        if u != v:\n",
        "            edge_list += [u,v]\n",
        "            adjacencies[u].add(v)\n",
        "            adjacencies[v].add(u)\n",
        "\n",
        "    edge_list += [PAD_TOKEN]*(SEQ_LEN-len(edge_list))\n",
        "    return edge_list, adjacencies\n",
        "\n",
        "\"\"\"\n",
        "input: G, represented as an adjacency list\n",
        "output: [INF]+[d(1,i) for i in range(n)] if target=None\n",
        "if target is set to some value, then we instead just output that specific distance\n",
        "\"\"\"\n",
        "def SSSP(G, target=None):\n",
        "    dist = [INF for _ in G]\n",
        "    dist[1] = 0\n",
        "    frontier = deque()\n",
        "    frontier.append(1)\n",
        "    while len(frontier) > 0:\n",
        "        vtx = frontier.popleft()\n",
        "        for x in G[vtx]:\n",
        "            if dist[x] == INF:\n",
        "                dist[x] = 1 + dist[vtx]\n",
        "                frontier.append(x)\n",
        "                if x == target:\n",
        "                    return dist[target]\n",
        "    if target is not None:\n",
        "        return dist[target]\n",
        "    else:\n",
        "        return dist\n",
        "\n",
        "graphs1 = []\n",
        "distance1 = []\n",
        "\n",
        "graphs2 = []\n",
        "distances2 = []\n",
        "\n",
        "for n in tqdm(range(8, MAX_VTXS)):\n",
        "    for _ in range(NTRAIN1//MAX_VTXS):\n",
        "        edge_list, adj_list = random_graph(n)\n",
        "        dist = SSSP(adj_list, target=2)\n",
        "\n",
        "        graphs1.append(edge_list)\n",
        "        distance1.append(dist)\n",
        "\n",
        "for n in range(8, MAX_VTXS//4):\n",
        "    for _ in range(NTRAIN2//MAX_VTXS):\n",
        "        edge_list, adj_list = random_graph(n)\n",
        "        distances = SSSP(adj_list)\n",
        "        graphs2.append(edge_list)\n",
        "        distances2.append(distances)\n",
        "\n",
        "split1 = int(len(graphs1)*3/4)\n",
        "split2 = int(len(graphs2)*3/4)\n",
        "\n",
        "data = {\n",
        "    \"train1-data\": graphs1[:split1],\n",
        "    \"train1-labels\": distance1[:split1],\n",
        "    \"test1-data\": graphs1[split1:],\n",
        "    \"test1-labels\": distance1[split1:],\n",
        "    \"train2-data\": graphs2[:split2],\n",
        "    \"train2-labels\": distances2[:split2],\n",
        "    \"test2-data\": graphs2[split2:],\n",
        "    \"test2-labels\": distances2[split2:]\n",
        "}\n",
        "\n",
        "with open('data.pkl', 'wb') as file:\n",
        "    pickle.dump(data, file)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 2: Define Transformer Model"
      ],
      "metadata": {
        "id": "Q3Cg_8UQep8g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerModel(nn.Module):\n",
        "    def __init__(self, input_dim, model_dim, output_dim, num_heads, num_layers, seq_len, device, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.embedding = nn.Embedding(input_dim, model_dim)\n",
        "        self.model_dim = model_dim\n",
        "        self.seq_len = seq_len\n",
        "        self.device = device\n",
        "\n",
        "        encoder_layers = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads, dropout=dropout, batch_first=True)\n",
        "        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers)\n",
        "        self.fc_out = nn.Linear(model_dim*seq_len, output_dim)\n",
        "\n",
        "    def positional_encoding(self, batch_size):\n",
        "        pos_encoding = torch.arange(self.seq_len, device=self.device).unsqueeze(1) % 2\n",
        "        pos_encoding = pos_encoding.float().unsqueeze(0).repeat(batch_size, 1, 1)\n",
        "        return pos_encoding\n",
        "\n",
        "    def forward(self, src, key_padding_mask):\n",
        "        batch_size, src_len = src.size(0), src.size(1)\n",
        "        src_pos = self.positional_encoding(batch_size)\n",
        "        embed = self.embedding(src)\n",
        "        src = embed * sqrt(self.model_dim) + src_pos\n",
        "\n",
        "        output = self.transformer_encoder(src, None, src_key_padding_mask=key_padding_mask)\n",
        "        flat_output = torch.flatten(output, start_dim=1, end_dim=2)\n",
        "        output = self.fc_out(flat_output)\n",
        "        return output\n"
      ],
      "metadata": {
        "id": "tLOWhg_CeWzH"
      },
      "execution_count": 99,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 3: Load Data"
      ],
      "metadata": {
        "id": "bpIeg86S-hBb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "assert device.type == 'cuda', \"CUDA is not available. Please check your GPU setup.\"\n",
        "\n",
        "# PARAMS\n",
        "VOCAB_SIZE = 64 # one more than the max number of vertices\n",
        "model_dim = 512 # Dimension of model (embedding and transformer)\n",
        "num_epochs = 4\n",
        "batch_size = 32\n",
        "learning_rate = 0.001\n",
        "max_seq_len = 128\n",
        "num_heads = 8\n",
        "num_layers = 6\n",
        "PAD_TOKEN = 0\n",
        "model = TransformerModel(input_dim=VOCAB_SIZE, model_dim=model_dim,\n",
        "                         output_dim=VOCAB_SIZE, num_heads=num_heads,\n",
        "                         num_layers=num_layers, seq_len=max_seq_len,\n",
        "                         device=device).to(device)\n",
        "\n",
        "with open(\"data.pkl\", \"rb\") as f:\n",
        "    data = pickle.load(f)\n",
        "\n",
        "train_data1 = data[\"train1-data\"]\n",
        "train_label1 = data[\"train1-labels\"]\n",
        "train_data_tensor = torch.tensor(train_data1, dtype=torch.long, device=device)\n",
        "train_label_tensor = torch.tensor(train_label1, dtype=torch.long, device=device)\n",
        "train_padding_mask = (train_data_tensor != PAD_TOKEN).bool().to(device)\n",
        "train_dataset = TensorDataset(train_data_tensor, train_label_tensor, train_padding_mask)\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "test_data1 = data[\"test1-data\"]\n",
        "test_label1 = data[\"test1-labels\"]\n",
        "test_data_tensor = torch.tensor(test_data1, dtype=torch.long, device=device)\n",
        "test_label_tensor = torch.tensor(test_label1, dtype=torch.long, device=device)\n",
        "test_padding_mask = (test_data_tensor != PAD_TOKEN).bool().to(device)\n",
        "test_dataset = TensorDataset(test_data_tensor, test_label_tensor, test_padding_mask)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "train_accuracy = []\n",
        "test_accuracy = []"
      ],
      "metadata": {
        "id": "kWXvJRDYgFVP"
      },
      "execution_count": 100,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 4: Train the Model for the first task"
      ],
      "metadata": {
        "id": "f8Zn33m7CxL5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(num_epochs):\n",
        "    model.train() # set to training mode\n",
        "    epoch_loss = 0\n",
        "    correct_train = 0\n",
        "    total_train = 0\n",
        "\n",
        "    for batch_src, batch_labels, batch_padding_mask in train_loader:\n",
        "        optimizer.zero_grad()\n",
        "        output = model(batch_src, batch_padding_mask)\n",
        "\n",
        "        _, predicted = torch.max(output, 1)\n",
        "        correct_train += (predicted == batch_labels).sum().item()\n",
        "        total_train += batch_labels.size(0)\n",
        "\n",
        "        loss = criterion(output, batch_labels)\n",
        "        epoch_loss += loss.item()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Evaluate performance\n",
        "    model.eval()\n",
        "    correct_test = 0\n",
        "    total_test = 0\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch_src, batch_labels, batch_padding_mask in test_loader:\n",
        "            output = model(batch_src, batch_padding_mask)\n",
        "\n",
        "            _, predicted = torch.max(output, 1)\n",
        "            correct_test += (predicted == batch_labels).sum().item()\n",
        "            total_test += batch_labels.size(0)\n",
        "\n",
        "    epoch_test_acc = correct_test / total_test\n",
        "    epoch_train_acc = correct_train / total_train\n",
        "    test_accuracy.append(epoch_test_acc)\n",
        "    train_accuracy.append(epoch_train_acc)\n",
        "    print(f\"Epoch {epoch + 1}/{num_epochs} \\t Train Accuracy: {epoch_train_acc:.4f} \\t Test Accuracy: {epoch_test_acc:.4f}\")\n",
        "\n",
        "plt.figure(figsize=(10, 5))\n",
        "plt.plot(test_accuracy, label='Test', color='red')\n",
        "plt.plot(train_accuracy, label='Train', color='red')\n",
        "plt.title('Accuracy vs Epochs')\n",
        "plt.xlabel('Epochs'); plt.ylabel('Accuracy')\n",
        "plt.legend(); plt.grid()\n",
        "plt.show()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pvTfzGmCeXU4",
        "outputId": "5231507f-7a52-4eb7-893a-c49ca93b8baf"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/4 \t Train Accuracy: 0.2067 \t Test Accuracy: 0.3838\n",
            "Epoch 2/4 \t Train Accuracy: 0.2457 \t Test Accuracy: 0.3838\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 5: Fine Tune"
      ],
      "metadata": {
        "id": "LC6Xv3YfC0Rm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 6: Test generalization"
      ],
      "metadata": {
        "id": "JtTLXn4zC1z_"
      }
    }
  ]
}