cs762/mirror.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c8f3b368-b56f-4e15-a5a8-5bb42fcfa283",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "import networkx as nx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1e98b589-ca11-493c-84bb-a09e7764a393",
   "metadata": {},
   "outputs": [],
   "source": [
    "graphs = []\n",
    "\n",
    "ngraphs = 10000\n",
    "for idx in range(ngraphs):\n",
    "    n = 10 # Graph size\n",
    "    g = nx.path_graph(n) # Path graph\n",
    "    edge_index = torch.tensor(list(g.edges())).t().contiguous()\n",
    "\n",
    "    # Data and labels are mirrors\n",
    "    x = torch.randperm(n)\n",
    "    y = torch.flip(x, [0])\n",
    "    x = x.to(dtype=torch.float32).unsqueeze(dim=1)\n",
    "\n",
    "    # Mask\n",
    "    k = n//10\n",
    "    mask = torch.cat((torch.ones(k), -torch.ones(k), torch.zeros(n - 2 * k)))\n",
    "    mask = mask[torch.randperm(n)]\n",
    "    train_mask = mask == 0\n",
    "    val_mask = mask == 1\n",
    "    test_mask = mask == -1\n",
    "\n",
    "    val_mask[(idx+1)%3] = 1\n",
    "    test_mask[(idx+2)%3] = 1\n",
    "    train_mask[val_mask | test_mask] = 0\n",
    "\n",
    "    # Create the graph\n",
    "    data = Data(edge_index=edge_index, x=x, y=y, train_mask=train_mask, val_mask=val_mask, test_mask=test_mask)\n",
    "    graphs.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "49bc2f6f-6ee0-4f89-9215-93d0fb4a615c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "fname = \"mirror.pkl\"\n",
    "\n",
    "with open(fname, \"wb\") as f:\n",
    "    pickle.dump(graphs, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed743508-8f69-404a-a22a-8bfb22ef684a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hashlib\n",
    "\n",
    "def calculate_md5(file_path):\n",
    "    # Create an MD5 hash object\n",
    "    md5 = hashlib.md5()\n",
    "\n",
    "    # Open the file in binary mode\n",
    "    with open(file_path, \"rb\") as file:\n",
    "        # Read the file in chunks\n",
    "        chunk_size = 8192\n",
    "        for chunk in iter(lambda: file.read(chunk_size), b\"\"):\n",
    "            # Update the hash with the current chunk\n",
    "            md5.update(chunk)\n",
    "\n",
    "    # Return the hexadecimal representation of the digest\n",
    "    return md5.hexdigest()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b694fc-7196-4891-98ed-201eaac6777c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import InMemoryDataset, download_url\n",
    "import os\n",
    "import os.path as osp\n",
    "import shutil\n",
    "import fsspec\n",
    "import io\n",
    "\n",
    "def torch_save(data, path) -> None:\n",
    "    buffer = io.BytesIO()\n",
    "    torch.save(data, buffer)\n",
    "    with fsspec.open(path, 'wb') as f:\n",
    "        f.write(buffer.getvalue())\n",
    "\n",
    "def torch_load(path):\n",
    "    with fsspec.open(path, 'rb') as f:\n",
    "        return torch.load(f)\n",
    "\n",
    "class Mirror(InMemoryDataset):\n",
    "    def __init__(self):\n",
    "        super().__init__('/tmp/Mirror')\n",
    "        path = osp.join(self.processed_dir, self.processed_file_names[0])\n",
    "        self.load(path)\n",
    "\n",
    "    def download(self):\n",
    "        src = \"/nobackup/vbalivada/GraphGPS/cs762/mirror.pkl\"\n",
    "        dst = osp.join(self.raw_dir, \"mirror.pkl\")\n",
    "        shutil.copy(src, dst)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return ['mirror.pkl']\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return ['data.pt']\n",
    "\n",
    "    def process(self):\n",
    "        with open(osp.join(self.raw_dir, \"mirror.pkl\"), \"rb\") as f:\n",
    "            graphs = pickle.load(f)\n",
    "        self.save(self.__class__, graphs, osp.join(self.processed_dir, self.processed_file_names[0]))\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return [\"data.pt\"]\n",
    "\n",
    "    @staticmethod\n",
    "    def save(cls, data_list, path):\n",
    "        r\"\"\"Saves a list of data objects to the file path :obj:`path`.\"\"\"\n",
    "        data, slices = cls.collate(data_list)\n",
    "        torch_save((data.to_dict(), slices), path)\n",
    "    \n",
    "    def load(self, path):\n",
    "        r\"\"\"Loads the dataset from the file path :obj:`path`.\"\"\"\n",
    "        data, self.slices = torch_load(path)\n",
    "        if isinstance(data, dict):  # Backward compatibility.\n",
    "            data = Data.from_dict(data)\n",
    "        self.data = data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efeca94a-d28c-46d4-9034-819dfd1b67be",
   "metadata": {},
   "outputs": [],
   "source": [
    "mirror_dataset = Mirror()\n",
    "g = mirror_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7bc9a82-896a-4e40-851c-426208ce2b9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "from tempfile import TemporaryDirectory\n",
    "from typing import Tuple\n",
    "\n",
    "import torch\n",
    "from torch import nn, Tensor\n",
    "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
    "from torch.utils.data import dataset\n",
    "\n",
    "class TransformerModel(nn.Module):\n",
    "\n",
    "    def __init__(self, ntoken: int, d_model: int, nhead: int, d_hid: int,\n",
    "                 nlayers: int, dropout: float = 0.5):\n",
    "        super().__init__()\n",
    "        self.model_type = 'Transformer'\n",
    "        self.pos_encoder = PositionalEncoding(d_model, dropout)\n",
    "        encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout)\n",
    "        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)\n",
    "        self.embedding = nn.Embedding(ntoken, d_model)\n",
    "        self.d_model = d_model\n",
    "        self.linear = nn.Linear(d_model, ntoken)\n",
    "\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self) -> None:\n",
    "        initrange = 0.1\n",
    "        self.embedding.weight.data.uniform_(-initrange, initrange)\n",
    "        self.linear.bias.data.zero_()\n",
    "        self.linear.weight.data.uniform_(-initrange, initrange)\n",
    "\n",
    "    def forward(self, src: Tensor, src_mask: Tensor = None) -> Tensor:\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            src: Tensor, shape ``[seq_len, batch_size]``\n",
    "            src_mask: Tensor, shape ``[seq_len, seq_len]``\n",
    "\n",
    "        Returns:\n",
    "            output Tensor of shape ``[seq_len, batch_size, ntoken]``\n",
    "        \"\"\"\n",
    "        src = self.embedding(src) * math.sqrt(self.d_model)\n",
    "        src = self.pos_encoder(src)\n",
    "        if src_mask is None:\n",
    "            \"\"\"Generate a square causal mask for the sequence. The masked positions are filled with float('-inf').\n",
    "            Unmasked positions are filled with float(0.0).\n",
    "            \"\"\"\n",
    "            src_mask = nn.Transformer.generate_square_subsequent_mask(len(src)).to(device)\n",
    "        output = self.transformer_encoder(src, src_mask)\n",
    "        output = self.linear(output)\n",
    "        return output\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "        position = torch.arange(max_len).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
    "        pe = torch.zeros(max_len, 1, d_model)\n",
    "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``\n",
    "        \"\"\"\n",
    "        x = x + self.pe[:x.size(0), 0, :]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58f8e565-9334-4934-afff-21ba890031e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# device = 'cpu'\n",
    "model = TransformerModel(ntoken=10, d_model=8, nhead=2, d_hid=8, nlayers=2, dropout=0.1).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
    "\n",
    "model.train()\n",
    "for epoch in range(200):\n",
    "    for data in mirror_dataset:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(data.x.squeeze().to(torch.long))\n",
    "        loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])\n",
    "        loss.backward()\n",
    "        optimizer.step()    \n",
    "    print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5adc43de-ab78-42cb-b099-fadf5d0ac83e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "pred = model(data).argmax(dim=1)\n",
    "correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()\n",
    "acc = int(correct) / int(data.test_mask.sum())\n",
    "print(f'Accuracy: {acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "550a2e2b-c0d8-4ab9-8425-8bc8575e5dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(dataset.num_node_features, 16)\n",
    "        self.conv2 = GCNConv(16, dataset.num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9126ddfc-ffac-4bf1-9bfb-383299a5494a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = GCN().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
    "\n",
    "model.train()\n",
    "for epoch in range(200):\n",
    "    for data in dataset:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(data)\n",
    "        loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7d76fc1-d538-4984-a529-5ceb0b1168ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "pred = model(data).argmax(dim=1)\n",
    "correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()\n",
    "acc = int(correct) / int(data.test_mask.sum())\n",
    "print(f'Accuracy: {acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9568184-f546-49d0-9d72-3a01129e179c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred[data.test_mask], data.y[data.test_mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f18fae7-3df1-4203-9fec-b494add66170",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}