examples/ogbn_mag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rrgcn import RRGCNEmbedder\n",
    "import torch\n",
    "from ogb.nodeproppred import PygNodePropPredDataset\n",
    "from ogb.nodeproppred import Evaluator\n",
    "from torch import nn\n",
    "from tqdm import tqdm\n",
    "import copy\n",
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install ogb\n",
    "dataset = PygNodePropPredDataset(name=\"ogbn-mag\")\n",
    "\n",
    "split_idx = dataset.get_idx_split()\n",
    "train_idx, valid_idx, test_idx = (\n",
    "    split_idx[\"train\"],\n",
    "    split_idx[\"valid\"],\n",
    "    split_idx[\"test\"],\n",
    ")\n",
    "graph = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(\n",
       "  num_nodes_dict={\n",
       "    author=1134649,\n",
       "    field_of_study=59965,\n",
       "    institution=8740,\n",
       "    paper=736389\n",
       "  },\n",
       "  edge_index_dict={\n",
       "    (author, affiliated_with, institution)=[2, 1043998],\n",
       "    (author, writes, paper)=[2, 7145660],\n",
       "    (paper, cites, paper)=[2, 5416271],\n",
       "    (paper, has_topic, field_of_study)=[2, 7505078]\n",
       "  },\n",
       "  x_dict={ paper=[736389, 128] },\n",
       "  node_year={ paper=[736389, 1] },\n",
       "  edge_reltype={\n",
       "    (author, affiliated_with, institution)=[1043998, 1],\n",
       "    (author, writes, paper)=[7145660, 1],\n",
       "    (paper, cites, paper)=[5416271, 1],\n",
       "    (paper, has_topic, field_of_study)=[7505078, 1]\n",
       "  },\n",
       "  y_dict={ paper=[736389, 1] }\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_indices = []\n",
    "edge_types = []\n",
    "\n",
    "node_types = set()\n",
    "for k, v in graph.edge_index_dict.items():\n",
    "    node_types = node_types.union({k[0], k[2]})\n",
    "\n",
    "node_types = sorted(node_types)\n",
    "num_node_types = len(node_types)\n",
    "node_type_to_add = {n: i for i, n in enumerate(node_types)}\n",
    "\n",
    "for k, v in graph.edge_index_dict.items():\n",
    "    edge_indices.append(\n",
    "        torch.vstack(\n",
    "            (\n",
    "                ((v[0] * num_node_types) + node_type_to_add[k[0]]),\n",
    "                ((v[1] * num_node_types) + node_type_to_add[k[2]]),\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    edge_types.append(graph.edge_reltype[k])\n",
    "\n",
    "edge_index = torch.hstack(edge_indices)\n",
    "edge_type = torch.vstack(edge_types).squeeze()\n",
    "num_nodes = sum(graph.num_nodes_dict.values())\n",
    "num_rels = edge_type.unique().numel()\n",
    "\n",
    "# convert node idx to consecutive integers\n",
    "node_idx = torch.full((edge_index.max() + 1,), -1)\n",
    "node_idx[edge_index.unique()] = torch.arange(num_nodes)\n",
    "edge_index = node_idx[edge_index]\n",
    "assert num_nodes == edge_index.max() + 1\n",
    "\n",
    "# inverses\n",
    "edge_type = torch.hstack((2 * edge_type, (2 * edge_type) + 1))\n",
    "edge_index = torch.hstack((edge_index, edge_index[[1, 0]]))\n",
    "\n",
    "node_features = {}\n",
    "for i, (k, word_feat) in enumerate(graph.x_dict.items()):\n",
    "    node_features[i] = [\n",
    "        node_idx[\n",
    "            (torch.arange(word_feat.shape[0]) * num_node_types) + node_type_to_add[k]\n",
    "        ],\n",
    "        word_feat\n",
    "    ]\n",
    "\n",
    "train_ys = []\n",
    "valid_ys = []\n",
    "test_ys = []\n",
    "\n",
    "for k, v in graph.y_dict.items():\n",
    "    train_ys.append(v[train_idx[k]])\n",
    "    valid_ys.append(v[valid_idx[k]])\n",
    "    test_ys.append(v[test_idx[k]])\n",
    "\n",
    "train_y = torch.hstack(train_ys)\n",
    "valid_y = torch.hstack(valid_ys)\n",
    "test_y = torch.hstack(test_ys)\n",
    "\n",
    "train_idxs = []\n",
    "valid_idxs = []\n",
    "test_idxs = []\n",
    "\n",
    "for k in train_idx.keys():\n",
    "    train_idxs.append(node_idx[(train_idx[k] * num_node_types) + node_type_to_add[k]])\n",
    "    valid_idxs.append(node_idx[(valid_idx[k] * num_node_types) + node_type_to_add[k]])\n",
    "    test_idxs.append(node_idx[(test_idx[k] * num_node_types) + node_type_to_add[k]])\n",
    "\n",
    "train_idx = torch.hstack(train_idxs)\n",
    "valid_idx = torch.hstack(valid_idxs)\n",
    "test_idx = torch.hstack(test_idxs)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "embedder = RRGCNEmbedder(\n",
    "    num_nodes=num_nodes,\n",
    "    num_relations=num_rels,\n",
    "    num_layers=2,\n",
    "    emb_size=750,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [01:23<00:00, 83.28s/it]\n"
     ]
    }
   ],
   "source": [
    "# for node features to work well, they have to be normalized\n",
    "# you can choose \"standard\" for StandardScaler, \"robust\" for RobustScaler, \"quantile\"\n",
    "# for QuantileTransformer and \"power\" for PowerTransformer\n",
    "#\n",
    "# you could also pass sklearn compatible scalers by passing a dict keyed by\n",
    "# literal type, e.g.:\n",
    "# {0: StandardScaler(), 1: RobustScaler()}\n",
    "\n",
    "train_embs = embedder.embeddings(\n",
    "    edge_index,\n",
    "    edge_type,\n",
    "    node_features=node_features,\n",
    "    node_features_scalers=\"standard\",\n",
    "    idx=train_idx,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:56<00:00, 56.14s/it]\n",
      "100%|██████████| 1/1 [00:52<00:00, 52.81s/it]\n"
     ]
    }
   ],
   "source": [
    "# only fit node feature scalers on nodes reachable from train nodes,\n",
    "# for val and test nodes, reuse the fit scalers using embedder.get_last_fit_scalers()\n",
    "val_embs = embedder.embeddings(\n",
    "    edge_index,\n",
    "    edge_type,\n",
    "    node_features=node_features,\n",
    "    node_features_scalers=embedder.get_last_fit_scalers(),\n",
    "    idx=valid_idx,\n",
    ")\n",
    "test_embs = embedder.embeddings(\n",
    "    edge_index,\n",
    "    edge_type,\n",
    "    node_features=node_features,\n",
    "    node_features_scalers=embedder.get_last_fit_scalers(),\n",
    "    idx=test_idx,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler()\n",
    "train_embs_scaled = torch.tensor(scaler.fit_transform(train_embs), dtype=torch.float32)\n",
    "test_embs_scaled = torch.tensor(scaler.transform(test_embs), dtype=torch.float32)\n",
    "val_embs_scaled = torch.tensor(scaler.transform(val_embs), dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epochs 1/100: 100%|██████████| 63/63 [00:14<00:00,  4.28it/s, total_loss=3.32, val_loss=2.65]\n",
      "Epochs 2/100: 100%|██████████| 63/63 [00:14<00:00,  4.40it/s, total_loss=2.47, val_loss=2.45]\n",
      "Epochs 3/100: 100%|██████████| 63/63 [00:14<00:00,  4.38it/s, total_loss=2.29, val_loss=2.38]\n",
      "Epochs 4/100: 100%|██████████| 63/63 [00:14<00:00,  4.34it/s, total_loss=2.19, val_loss=2.29]\n",
      "Epochs 5/100: 100%|██████████| 63/63 [00:14<00:00,  4.25it/s, total_loss=2.12, val_loss=2.25]\n",
      "Epochs 6/100: 100%|██████████| 63/63 [00:14<00:00,  4.31it/s, total_loss=2.05, val_loss=2.23]\n",
      "Epochs 7/100: 100%|██████████| 63/63 [00:14<00:00,  4.36it/s, total_loss=2, val_loss=2.2]\n",
      "Epochs 8/100: 100%|██████████| 63/63 [00:14<00:00,  4.36it/s, total_loss=1.96, val_loss=2.19]\n",
      "Epochs 9/100: 100%|██████████| 63/63 [00:14<00:00,  4.22it/s, total_loss=1.92, val_loss=2.18]\n",
      "Epochs 10/100: 100%|██████████| 63/63 [00:14<00:00,  4.23it/s, total_loss=1.88, val_loss=2.16]\n",
      "Epochs 11/100: 100%|██████████| 63/63 [00:14<00:00,  4.35it/s, total_loss=1.85, val_loss=2.17]\n",
      "Epochs 12/100: 100%|██████████| 63/63 [00:14<00:00,  4.43it/s, total_loss=1.82, val_loss=2.16]\n",
      "Epochs 13/100: 100%|██████████| 63/63 [00:13<00:00,  4.66it/s, total_loss=1.79, val_loss=2.13]\n",
      "Epochs 14/100: 100%|██████████| 63/63 [00:14<00:00,  4.45it/s, total_loss=1.76, val_loss=2.12]\n",
      "Epochs 15/100: 100%|██████████| 63/63 [00:13<00:00,  4.68it/s, total_loss=1.73, val_loss=2.11]\n",
      "Epochs 16/100: 100%|██████████| 63/63 [00:14<00:00,  4.47it/s, total_loss=1.71, val_loss=2.11]\n",
      "Epochs 17/100: 100%|██████████| 63/63 [00:13<00:00,  4.73it/s, total_loss=1.69, val_loss=2.12]\n",
      "Epochs 18/100: 100%|██████████| 63/63 [00:14<00:00,  4.31it/s, total_loss=1.67, val_loss=2.1]\n",
      "Epochs 19/100: 100%|██████████| 63/63 [00:13<00:00,  4.56it/s, total_loss=1.64, val_loss=2.14]\n",
      "Epochs 20/100: 100%|██████████| 63/63 [00:14<00:00,  4.36it/s, total_loss=1.63, val_loss=2.12]\n",
      "Epochs 21/100: 100%|██████████| 63/63 [00:14<00:00,  4.45it/s, total_loss=1.61, val_loss=2.16]\n",
      "Epochs 22/100: 100%|██████████| 63/63 [00:14<00:00,  4.29it/s, total_loss=1.58, val_loss=2.12]\n",
      "Epochs 23/100: 100%|██████████| 63/63 [00:13<00:00,  4.58it/s, total_loss=1.57, val_loss=2.12]\n",
      "Epochs 24/100: 100%|██████████| 63/63 [00:14<00:00,  4.26it/s, total_loss=1.55, val_loss=2.13]\n",
      "Epochs 25/100: 100%|██████████| 63/63 [00:13<00:00,  4.54it/s, total_loss=1.53, val_loss=2.12]\n",
      "Epochs 26/100: 100%|██████████| 63/63 [00:14<00:00,  4.35it/s, total_loss=1.52, val_loss=2.13]\n",
      "Epochs 27/100: 100%|██████████| 63/63 [00:13<00:00,  4.52it/s, total_loss=1.5, val_loss=2.13]\n",
      "Epochs 28/100: 100%|██████████| 63/63 [00:14<00:00,  4.29it/s, total_loss=1.49, val_loss=2.15]\n",
      "Epochs 29/100: 100%|██████████| 63/63 [00:13<00:00,  4.61it/s, total_loss=1.47, val_loss=2.17]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 00029: reducing learning rate of group 0 to 1.0000e-04.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epochs 30/100: 100%|██████████| 63/63 [00:13<00:00,  4.51it/s, total_loss=1.38, val_loss=2.14]\n",
      "Epochs 31/100: 100%|██████████| 63/63 [00:13<00:00,  4.69it/s, total_loss=1.35, val_loss=2.14]\n",
      "Epochs 32/100: 100%|██████████| 63/63 [00:14<00:00,  4.27it/s, total_loss=1.34, val_loss=2.14]\n",
      "Epochs 33/100: 100%|██████████| 63/63 [00:13<00:00,  4.52it/s, total_loss=1.34, val_loss=2.14]\n",
      "Epochs 34/100: 100%|██████████| 63/63 [00:15<00:00,  4.19it/s, total_loss=1.33, val_loss=2.14]\n",
      "Epochs 35/100: 100%|██████████| 63/63 [00:14<00:00,  4.22it/s, total_loss=1.32, val_loss=2.15]\n",
      "Epochs 36/100: 100%|██████████| 63/63 [00:14<00:00,  4.31it/s, total_loss=1.32, val_loss=2.15]\n",
      "Epochs 37/100: 100%|██████████| 63/63 [00:13<00:00,  4.61it/s, total_loss=1.32, val_loss=2.14]\n",
      "Epochs 38/100: 100%|██████████| 63/63 [00:14<00:00,  4.26it/s, total_loss=1.31, val_loss=2.15]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Early stopping, resetting weights to best modelwith val_loss 2.0962438583374023\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# batched training is a mess in CatBoost, so use MLP for this trainset\n",
    "train_set = torch.utils.data.TensorDataset(train_embs_scaled, train_y)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=10_000, shuffle=True)\n",
    "\n",
    "mlp = nn.Sequential(\n",
    "    nn.Linear(train_embs.shape[1], train_embs.shape[1] // 2),\n",
    "    nn.Dropout(0.3),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(train_embs.shape[1] // 2, train_embs.shape[1] // 2),\n",
    "    nn.Dropout(0.1),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(train_embs.shape[1] // 2, train_y.unique().numel() * 2),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(train_y.unique().numel() * 2, train_y.unique().numel()),\n",
    ")\n",
    "mlp = mlp.to(device)\n",
    "\n",
    "loss = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-3)\n",
    "num_epochs = 100\n",
    "early_stopping_epochs = 20\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer, patience=early_stopping_epochs // 2,\n",
    "    verbose=True\n",
    ")\n",
    "\n",
    "best_val_loss = 1_000_000\n",
    "best_model = None\n",
    "epochs_since_best = 0\n",
    "for epoch in range(num_epochs):\n",
    "    mlp.train()\n",
    "    with tqdm(\n",
    "        train_loader, total=len(train_loader), desc=f\"Epochs {epoch + 1}/{num_epochs}\"\n",
    "    ) as bar:\n",
    "        total_loss = 0\n",
    "        total_items = 0\n",
    "        for i, (x, y) in enumerate(bar):\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            out = mlp(x.to(device))\n",
    "            l = loss(out, y.squeeze().to(device))\n",
    "            total_loss += (x.shape[0]) * l.item()\n",
    "            total_items += x.shape[0]\n",
    "\n",
    "            l.backward()\n",
    "            optimizer.step()\n",
    "            bar.set_postfix(loss=l.item())\n",
    "\n",
    "            if i == len(train_loader) - 1:\n",
    "                mlp.eval()\n",
    "                val_loss = loss(\n",
    "                    mlp(val_embs_scaled.to(device)), valid_y.squeeze().to(device)\n",
    "                )\n",
    "\n",
    "                if val_loss.item() < best_val_loss:\n",
    "                    best_model = copy.deepcopy(mlp)\n",
    "                    best_val_loss = val_loss.item()\n",
    "                    epochs_since_best = 0\n",
    "                else:\n",
    "                    epochs_since_best += 1\n",
    "\n",
    "                bar.set_postfix(\n",
    "                    total_loss=total_loss / total_items, val_loss=val_loss.item()\n",
    "                )\n",
    "                scheduler.step(val_loss)\n",
    "\n",
    "    if epochs_since_best == early_stopping_epochs:\n",
    "        mlp = best_model\n",
    "        print(\n",
    "            \"Early stopping, resetting weights to best model\"\n",
    "            + f\"with val_loss {best_val_loss}\"\n",
    "        )\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'acc': 0.396027563842724}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlp.eval()\n",
    "evaluator = Evaluator(name=\"ogbn-mag\")\n",
    "\n",
    "print(\"Test:\")\n",
    "evaluator.eval(\n",
    "    {\n",
    "        \"y_true\": test_y.cpu().numpy(),\n",
    "        \"y_pred\": mlp(test_embs_scaled.to(device))\n",
    "        .argmax(-1)\n",
    "        .reshape(-1, 1)\n",
    "        .detach()\n",
    "        .cpu(),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'acc': 0.4113811865164383}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Validation:\")\n",
    "evaluator.eval(\n",
    "    {\n",
    "        \"y_true\": valid_y.cpu().numpy(),\n",
    "        \"y_pred\": mlp(val_embs_scaled.to(device))\n",
    "        .argmax(-1)\n",
    "        .reshape(-1, 1)\n",
    "        .detach()\n",
    "        .cpu(),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trained full-batch R-GCN with input node features on the same dataset\n",
    "# https://github.com/snap-stanford/ogb/blob/master/examples/nodeproppred/mag/rgcn.py\n",
    "# https://ogb.stanford.edu/docs/leader_nodeprop/#ogbn-mag\n",
    "# -------------------------------------\n",
    "#   Test accuracy     Valid accuracy \n",
    "# -------------------------------------\n",
    "#   0.3977 ± 0.0046   0.4084 ± 0.0041\n",
    "# -------------------------------------"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
  },
  "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}