pytorch_lightgcn_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d03d1d50-0dff-4e2b-a55c-258ceb0d9430",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T17:57:49.598875Z",
     "iopub.status.busy": "2023-08-16T17:57:49.598781Z",
     "iopub.status.idle": "2023-08-16T17:57:50.648797Z",
     "shell.execute_reply": "2023-08-16T17:57:50.648431Z",
     "shell.execute_reply.started": "2023-08-16T17:57:49.598863Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch_geometric as PyG\n",
    "from torch_geometric.datasets import AmazonBook\n",
    "from torch_geometric.nn import LightGCN\n",
    "from torch_geometric.utils import degree\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8db3b34-3e60-4814-a959-fda64641ed78",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T17:57:50.649714Z",
     "iopub.status.busy": "2023-08-16T17:57:50.649508Z",
     "iopub.status.idle": "2023-08-16T17:57:50.728560Z",
     "shell.execute_reply": "2023-08-16T17:57:50.728268Z",
     "shell.execute_reply.started": "2023-08-16T17:57:50.649704Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HeteroData(\n",
      "  user={ num_nodes=52643 },\n",
      "  book={ num_nodes=91599 },\n",
      "  (user, rates, book)={\n",
      "    edge_index=[2, 2380730],\n",
      "    edge_label_index=[2, 603378],\n",
      "  },\n",
      "  (book, rated_by, user)={ edge_index=[2, 2380730] }\n",
      ")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 4761460], edge_label_index=[2, 603378], node_type=[144242], edge_type=[4761460])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = osp.join(osp.dirname('..'), 'data', 'Amazon')\n",
    "dataset = AmazonBook(path)\n",
    "data = dataset[0]\n",
    "print(data)\n",
    "num_users, num_books = data['user'].num_nodes, data['book'].num_nodes\n",
    "heterodata = data\n",
    "data = heterodata.to_homogeneous().to(device)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b1ff5603-d26a-4f38-bee0-b1e5c2fda077",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T17:57:50.729126Z",
     "iopub.status.busy": "2023-08-16T17:57:50.728993Z",
     "iopub.status.idle": "2023-08-16T17:57:50.734624Z",
     "shell.execute_reply": "2023-08-16T17:57:50.734340Z",
     "shell.execute_reply.started": "2023-08-16T17:57:50.729116Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total nodes: 144242\n"
     ]
    }
   ],
   "source": [
    "assert (heterodata['user', 'rates', 'book'].edge_index[0] < heterodata['user'].num_nodes).all()\n",
    "assert (heterodata['book', 'rated_by', 'user'].edge_index[1] < heterodata['book'].num_nodes).all()\n",
    "print(\"Total nodes:\", heterodata.num_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "733ce577-97f3-4537-b269-7752e90933d7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T17:57:50.735134Z",
     "iopub.status.busy": "2023-08-16T17:57:50.735030Z",
     "iopub.status.idle": "2023-08-16T18:00:07.725144Z",
     "shell.execute_reply": "2023-08-16T18:00:07.724745Z",
     "shell.execute_reply.started": "2023-08-16T17:57:50.735124Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. min: tensor(0) max: tensor(52642)\n",
      "1. min: tensor(0) max: tensor(91598)\n",
      "\n",
      "0. min: tensor(0) max: tensor(144241)\n",
      "1. min: tensor(0) max: tensor(144241)\n",
      "\n",
      "0. min: tensor(0) max: tensor(52642)\n",
      "1. min: tensor(52643) max: tensor(144241)\n",
      "\n",
      "0. min: tensor(52643) max: tensor(144241)\n",
      "1. min: tensor(0) max: tensor(52642)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# book -> user\n",
    "def _min_max(x):\n",
    "    print('0. min:', min(x[0]), 'max:', max(x[0]))\n",
    "    print('1. min:', min(x[1]), 'max:', max(x[1]))\n",
    "    print()\n",
    "\n",
    "_min_max(heterodata['user', 'rates', 'book'].edge_index)\n",
    "_min_max(data.edge_index)\n",
    "_min_max(data.edge_index[:, :2380730])\n",
    "_min_max(data.edge_index[:, 2380730:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b2ef72d7-000f-4fd4-b145-cea76c8ffeb7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:07.725664Z",
     "iopub.status.busy": "2023-08-16T18:00:07.725563Z",
     "iopub.status.idle": "2023-08-16T18:00:07.781672Z",
     "shell.execute_reply": "2023-08-16T18:00:07.781328Z",
     "shell.execute_reply.started": "2023-08-16T18:00:07.725654Z"
    }
   },
   "outputs": [],
   "source": [
    "# Use all message passing edges as training labels:\n",
    "batch_size = 8192\n",
    "# Convert to onedirectional edge\n",
    "mask = data.edge_index[0] < data.edge_index[1]\n",
    "train_edge_label_index = data.edge_index[:, mask]\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    range(train_edge_label_index.size(1)),\n",
    "    shuffle=True,\n",
    "    batch_size=batch_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b2737091-c168-412d-a0cc-b72a623ac120",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:07.782270Z",
     "iopub.status.busy": "2023-08-16T18:00:07.782154Z",
     "iopub.status.idle": "2023-08-16T18:00:35.827226Z",
     "shell.execute_reply": "2023-08-16T18:00:35.826916Z",
     "shell.execute_reply.started": "2023-08-16T18:00:07.782259Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. min: tensor(0) max: tensor(52642)\n",
      "1. min: tensor(52643) max: tensor(144241)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "_min_max(train_edge_label_index)\n",
    "assert (data.edge_index[:, data.edge_type == 0] == train_edge_label_index).all()\n",
    "assert (train_edge_label_index[0] < 52643).all()\n",
    "assert (train_edge_label_index[1] >= 52643).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "861f6cc3-ebb7-4eec-bdbc-796c5829c1ef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:35.828664Z",
     "iopub.status.busy": "2023-08-16T18:00:35.828547Z",
     "iopub.status.idle": "2023-08-16T18:00:35.906767Z",
     "shell.execute_reply": "2023-08-16T18:00:35.906302Z",
     "shell.execute_reply.started": "2023-08-16T18:00:35.828654Z"
    }
   },
   "outputs": [],
   "source": [
    "model = LightGCN(\n",
    "    num_nodes=data.num_nodes,\n",
    "    embedding_dim=64,\n",
    "    num_layers=3,\n",
    ").to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a0accb1d-17e5-4e27-81ba-3300bfa1261f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:35.907323Z",
     "iopub.status.busy": "2023-08-16T18:00:35.907201Z",
     "iopub.status.idle": "2023-08-16T18:00:35.969208Z",
     "shell.execute_reply": "2023-08-16T18:00:35.968697Z",
     "shell.execute_reply.started": "2023-08-16T18:00:35.907311Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------------+-------------------------+----------------+-----------+\n",
      "| Layer                  | Input Shape             | Output Shape   | #Param    |\n",
      "|------------------------+-------------------------+----------------+-----------|\n",
      "| LightGCN               | [2, 8192]               | [8192]         | 9,231,488 |\n",
      "| ├─(embedding)Embedding | --                      | --             | 9,231,488 |\n",
      "| ├─(convs)ModuleList    | --                      | --             | --        |\n",
      "| │    └─(0)LGConv       | [144242, 64], [2, 8192] | [144242, 64]   | --        |\n",
      "| │    └─(1)LGConv       | [144242, 64], [2, 8192] | [144242, 64]   | --        |\n",
      "| │    └─(2)LGConv       | [144242, 64], [2, 8192] | [144242, 64]   | --        |\n",
      "+------------------------+-------------------------+----------------+-----------+\n"
     ]
    }
   ],
   "source": [
    "print(PyG.nn.summary(model, data.edge_index[:, :batch_size]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "76dc601b-77bc-4b5d-afc0-296d641c6bf2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:35.970061Z",
     "iopub.status.busy": "2023-08-16T18:00:35.969915Z",
     "iopub.status.idle": "2023-08-16T18:00:35.973503Z",
     "shell.execute_reply": "2023-08-16T18:00:35.973238Z",
     "shell.execute_reply.started": "2023-08-16T18:00:35.970049Z"
    }
   },
   "outputs": [],
   "source": [
    "def train():\n",
    "    total_loss = total_examples = 0\n",
    "\n",
    "    for index in tqdm(train_loader):\n",
    "        # Sample positive and negative labels.\n",
    "        pos_edge_label_index = train_edge_label_index[:, index]\n",
    "        neg_edge_label_index = torch.stack([\n",
    "            pos_edge_label_index[0],\n",
    "            torch.randint(num_users, num_users + num_books,\n",
    "                          (index.numel(), ), device=device)\n",
    "        ], dim=0)\n",
    "        edge_label_index = torch.cat([\n",
    "            pos_edge_label_index,\n",
    "            neg_edge_label_index,\n",
    "        ], dim=1)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # WHY DOES IT USE data.edge_index here??\n",
    "        # R: Because it needs return edges for message passing\n",
    "        # but don't worry, it should be homogeneized\n",
    "        pos_rank, neg_rank = model(data.edge_index, edge_label_index).chunk(2)\n",
    "\n",
    "        loss = model.recommendation_loss(\n",
    "            pos_rank,\n",
    "            neg_rank,\n",
    "            node_id=edge_label_index.unique(),\n",
    "        )\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += float(loss) * pos_rank.numel()\n",
    "        total_examples += pos_rank.numel()\n",
    "\n",
    "    return total_loss / total_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0d21cbeb-efb3-4936-9214-d3631cf3c920",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:35.973973Z",
     "iopub.status.busy": "2023-08-16T18:00:35.973878Z",
     "iopub.status.idle": "2023-08-16T18:00:47.686401Z",
     "shell.execute_reply": "2023-08-16T18:00:47.685720Z",
     "shell.execute_reply.started": "2023-08-16T18:00:35.973964Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(False)\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n"
     ]
    }
   ],
   "source": [
    "@torch.no_grad()\n",
    "def test(k: int):\n",
    "    # WHY DOES IT USE DATA.EDGE INDEX NOW??? AND NOT \n",
    "    emb = model.get_embedding(data.edge_index)\n",
    "    user_emb, book_emb = emb[:num_users], emb[num_users:]\n",
    "\n",
    "    print((model.convs[0](model.embedding.weight, data.edge_index) == 0).all())\n",
    "\n",
    "    precision = recall = total_examples = 0\n",
    "    for start in range(0, num_users, batch_size):\n",
    "        end = start + batch_size\n",
    "        logits = user_emb[start:end] @ book_emb.t()\n",
    "\n",
    "        # print('start:', start, 'end:', end)\n",
    "        # print('logits.size()', logits.size())\n",
    "\n",
    "        # Exclude training edges:\n",
    "        mask = ((train_edge_label_index[0] >= start) &\n",
    "                (train_edge_label_index[0] < end))\n",
    "        logits[train_edge_label_index[0, mask] - start,\n",
    "               train_edge_label_index[1, mask] - num_users] = float('-inf')\n",
    "\n",
    "        # print('mask.size()', mask.size())\n",
    "\n",
    "        # Computing precision and recall:\n",
    "        ground_truth = torch.zeros_like(logits, dtype=torch.bool)\n",
    "        mask = ((data.edge_label_index[0] >= start) &\n",
    "                (data.edge_label_index[0] < end))\n",
    "        ground_truth[data.edge_label_index[0, mask] - start,\n",
    "                     data.edge_label_index[1, mask] - num_users] = True\n",
    "        node_count = degree(data.edge_label_index[0, mask] - start,\n",
    "                            num_nodes=logits.size(0))\n",
    "\n",
    "        # print('ground_truth.size():', ground_truth.size())\n",
    "        # print('node_count.size():', node_count.size())\n",
    "\n",
    "        topk_index = logits.topk(k, dim=-1).indices\n",
    "        # print('topk_index.size()', topk_index.size())\n",
    "        isin_mat = ground_truth.gather(1, topk_index)\n",
    "\n",
    "        precision += float((isin_mat.sum(dim=-1) / k).sum())\n",
    "        recall += float((isin_mat.sum(dim=-1) / node_count.clamp(1e-6)).sum())\n",
    "        total_examples += int((node_count > 0).sum())\n",
    "        print(20*'=')\n",
    "\n",
    "    return precision / total_examples, recall / total_examples\n",
    "\n",
    "precision, recall = test(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9418e6-df4f-4970-81df-d0b8e5259cf8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-16T18:00:47.687473Z",
     "iopub.status.busy": "2023-08-16T18:00:47.687201Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 291/291 [24:21<00:00,  5.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(False)\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "====================\n",
      "Epoch: 001, Loss: 0.5014, Precision@20: 0.0049, Recall@20: 0.0100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                          | 0/291 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, 5):\n",
    "    loss = train()\n",
    "    precision, recall = test(k=20)\n",
    "    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Precision@20: '\n",
    "          f'{precision:.4f}, Recall@20: {recall:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b860bff-3bd5-4dfc-a286-5c9cc2fcdc61",
   "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.11.3"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}