lab07_recsys/lab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9A92ddPbyodS"
   },
   "source": [
    "# Lab — Graph Neural Networks for Recommender Systems\n",
    "\n",
    "Recommender system task can be formulated in terms of link prediction problem.\n",
    "\n",
    "\n",
    "Today we will work with the Movielens dataset. It contains ratings assigned to movie by user, movie genres and user features. The goal of the dataset is to predict the final rating. We reduce task to the prediction of interaction to fit the link prediction task.\n",
    "\n",
    "The assignment will consists of three tasks\n",
    "\n",
    "1. We will train simple GraphSAGE model with user and item encoders to predict the new links\n",
    "2. We will enhance the loss function with low-rank positive penalties following the Uber Eats idea\n",
    "3. We will add importances to the sampling strategy following the PinSAGE model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Dx_6A_R7hDW1",
    "outputId": "8e4be60c-2179-4840-99a7-c62a64f7c514"
   },
   "outputs": [],
   "source": [
    "! pip install dgl-cu111 -f https://data.dgl.ai/wheels/repo.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KCik0ClvDlPw"
   },
   "source": [
    "## Data preparation\n",
    "\n",
    "Before we start any model, we need to prepare datasets and define all required methods.\n",
    "\n",
    "Let us download the data from official repository. We will take pretty small data with only one million ratings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XK-D2RVmDm0o",
    "outputId": "8f1f746b-4608-4beb-a2b5-573d2d41775c"
   },
   "outputs": [],
   "source": [
    "! wget http://files.grouplens.org/datasets/movielens/ml-1m.zip\n",
    "! unzip ml-1m.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cc-5360QdutH"
   },
   "source": [
    "Let us see how data looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 261
    },
    "id": "ncRAd8joIBg0",
    "outputId": "f1c1f897-2e51-4334-c89d-cca6b799b0a4"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "ratings = pd.read_csv(\"ml-1m/ratings.dat\", sep=\"::\", header=None)\n",
    "ratings.columns = [\"user_id\", \"item_id\", \"rating\", \"timestamp\"]\n",
    "ratings.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r-Un-ztrdzmb"
   },
   "source": [
    "Let us check misses in the `user_id` column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "J3RO13avd0CI",
    "outputId": "e2374651-b3e2-441c-83f8-f616773066c3"
   },
   "outputs": [],
   "source": [
    "set(range(ratings.user_id.max() + 1)) - set(ratings.user_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q5QhrGg9dzZk"
   },
   "source": [
    "Looks ok, we need to substract `1` from `user_id` column for proper indexing of feature tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iGtgr6LUeVh7"
   },
   "outputs": [],
   "source": [
    "ratings.user_id -= 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BMWovcZjebl-"
   },
   "source": [
    "What about items?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 261
    },
    "id": "U55AWgHNJivB",
    "outputId": "bde3abff-bd0c-4cd8-f3c8-5a7b6810556f"
   },
   "outputs": [],
   "source": [
    "items = pd.read_csv(\"ml-1m/movies.dat\", sep=\"::\", encoding='cp1251', header=None)\n",
    "items.columns = [\"item_id\", \"title\", \"tags\"]\n",
    "items.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iHXGEBiUenpj",
    "outputId": "7e3ab084-4ed8-457f-8333-d082c75ed777"
   },
   "outputs": [],
   "source": [
    "len(set(range(items.item_id.max() + 1)) - set(items.item_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_R9-s-DOeuNB"
   },
   "source": [
    "Ok, there are missed item indices in ordering, so let us rearrange ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "71i8pcJ9enmy"
   },
   "outputs": [],
   "source": [
    "item_mapper = dict(zip(list(set(items.item_id)), range(items.item_id.nunique())))\n",
    "items.item_id = items.item_id.map(item_mapper)\n",
    "ratings.item_id = ratings.item_id.map(item_mapper)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UzOUCEeNfS6P"
   },
   "source": [
    "Movies are described by the set of tags, so we will use one-hot encoding as item features. Our goal is to build the inductive model that are suitable to handle cold-start problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "D8vl-HZLfsHI",
    "outputId": "39bff894-84ec-498c-c5a3-00aede518688"
   },
   "outputs": [],
   "source": [
    "item_features = pd.DataFrame(items.tags.str.split(\"|\").map(lambda x: {i: 1 for i in x}).tolist()).fillna(0).values\n",
    "item_features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z3BwubNAf3LO"
   },
   "source": [
    "Now, we are ready to work with the user features (do not forget to substract `1` from id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 261
    },
    "id": "Hk4qLIhrJjSU",
    "outputId": "0807da2e-a2bd-4e4c-d716-36e0fedd54c4"
   },
   "outputs": [],
   "source": [
    "users = pd.read_csv(\"ml-1m/users.dat\", sep=\"::\", header=None)\n",
    "users.columns = [\"user_id\", \"gender\", \"age\", \"occupation\", \"zipcode\"]\n",
    "users.user_id -= 1\n",
    "users.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rycy2TZogFzA"
   },
   "source": [
    "Age and occupation are defined by the groups, so it is better to encode them as one-hot vectors. Zipcode is very variative categorical, so we will just drop it for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YVpvauevgfRf"
   },
   "outputs": [],
   "source": [
    "# change gender to indicator\n",
    "users.gender = (users.gender == \"F\").astype(int)\n",
    "# extract OHE vector from age\n",
    "users = users.join(pd.get_dummies(users.age))\n",
    "# extract OHE vector from occupation\n",
    "users = users.join(pd.get_dummies(users.occupation), rsuffix=\"_occ\")\n",
    "# drop non-required fields\n",
    "user_features = users.drop([\"user_id\", \"age\", \"occupation\", \"zipcode\"], axis=1).values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xr-fiVSggqa3"
   },
   "source": [
    "Now, we are ready to define a graph in dgl format and assign features to nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uqLKeqvCQ2hP",
    "outputId": "bb001e7d-2b23-417d-e75f-822323415ff7"
   },
   "outputs": [],
   "source": [
    "import dgl\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# create bipartite graph for user-item interactions\n",
    "graph = dgl.graph((ratings.user_id, ratings.item_id + ratings.user_id.max() + 1))\n",
    "graph.edata[\"rating\"] = torch.Tensor(ratings.rating.values)\n",
    "# add reverse edges for message passing\n",
    "graph = dgl.add_reverse_edges(graph, copy_edata=True)\n",
    "# preserve type of item\n",
    "graph.ndata[\"ntype\"] = torch.cat((torch.zeros(ratings.user_id.max() + 1), torch.ones(ratings.item_id.max() + 1))).reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AureFdNHv0Jn"
   },
   "source": [
    "To simplify message passing code, we allign the user and item features to store it in the same graph attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ahClPy-wv7lI",
    "outputId": "0bfb3352-f841-499c-c4fd-29a49ec565ec"
   },
   "outputs": [],
   "source": [
    "print(user_features.shape, item_features.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zz3az2fEvzG0"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "item_features = np.hstack([item_features, np.zeros((item_features.shape[0], 11))])\n",
    "graph.ndata[\"feat\"] = torch.Tensor(np.vstack([user_features, item_features]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aN0lvEarswBE"
   },
   "source": [
    "We also need to split our graph on train, validation and test parts. We will split it by unique nodes, because GraphSAGE model is suitable to infer in inductive fashion (for previously unseen nodes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lnUEuTeRikR6"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def prepare_graph_for_loaders(graph, perm, l=None, r=None, cum=True):\n",
    "  real_nodes = perm[int((l or 0) * perm.shape[0]): int((r or 1) * perm.shape[0])]\n",
    "  g = graph.subgraph(perm[:int((r or 1) * perm.shape[0])] if cum else real_nodes)\n",
    "  n = torch.where(torch.isin(g.ndata[dgl.NID], torch.Tensor(real_nodes)))[0]\n",
    "  f, t = g.edges()\n",
    "  edge_filter = torch.logical_or(torch.isin(f, n), torch.isin(t, n))\n",
    "  return g, real_nodes, g.edges(\"eid\")[edge_filter]\n",
    "\n",
    "np.random.seed(0)\n",
    "perm = np.random.permutation(graph.nodes().numpy())\n",
    "\n",
    "train_graph, train_nodes, train_edges = prepare_graph_for_loaders(graph, perm, r=0.8)\n",
    "val_graph, val_nodes, val_edges = prepare_graph_for_loaders(graph, perm, l=0.8, r=0.9)\n",
    "test_graph, test_nodes, test_edges = prepare_graph_for_loaders(graph, perm, l=0.9, r=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ql8m4gzBmPFL"
   },
   "source": [
    "## Model definition\n",
    "\n",
    "User and item data has different types of features, so we need to project it in one space. We can use simple MLP encoders.\n",
    "\n",
    "The model will be a GraphSAGE. It works as follows:\n",
    "\n",
    "1. Sample several layers of neighbors uniformly\n",
    "2. Pass features through linear transformation\n",
    "3. Aggregate the features of adjacent nodes from neighbor layer\n",
    "4. Apply non-linearity\n",
    "5. Iterate from the farest neighbors to the current node\n",
    "\n",
    "Here we write the custom GraphSAGE because our graph is bipartite."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "popfXoC8RLwK"
   },
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "class GraphSAGE(nn.Module):\n",
    "    def __init__(self, n_input, n_hidden, n_layers, activation, dropout):\n",
    "        super(GraphSAGE, self).__init__()\n",
    "        self.n_layers = n_layers\n",
    "        self.n_hidden = n_hidden\n",
    "        self.layers = nn.ModuleList()\n",
    "\n",
    "        self.user_encoder = nn.Sequential(\n",
    "            nn.Linear(n_input, n_hidden * 2),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(n_hidden * 2, n_hidden),\n",
    "        )\n",
    "        self.item_encoder = nn.Sequential(\n",
    "            nn.Linear(n_input, n_hidden * 2),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(n_hidden * 2, n_hidden),\n",
    "        )\n",
    "\n",
    "        for i in range(n_layers):\n",
    "            self.layers.append(dgl.nn.SAGEConv(n_hidden, n_hidden, \"mean\"))\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, pos_graph, neg_graph, blocks):\n",
    "      u, i = self.user_encoder(blocks[0].ndata[\"feat\"][\"_N\"]), self.item_encoder(blocks[0].ndata[\"feat\"][\"_N\"])\n",
    "      h = (1 - blocks[0].ndata[\"ntype\"][\"_N\"]) * u + i * (blocks[0].ndata[\"ntype\"][\"_N\"])\n",
    "      h = F.normalize(h, dim=1)\n",
    "\n",
    "      for l, (layer, block) in enumerate(zip(self.layers, blocks)):\n",
    "          h = self.activation(h)\n",
    "          h = self.dropout(h)\n",
    "          h = layer(block, h, block.edata.get(\"weight\"))\n",
    "          h = F.normalize(h, dim=1)\n",
    "\n",
    "      s, d = pos_graph.edges()\n",
    "      pos_scores = (h[s] * h[d]).sum(dim=1)\n",
    "      if neg_graph is not None:\n",
    "          s, d = neg_graph.edges()\n",
    "          neg_scores = (h[s] * h[d]).sum(dim=1)\n",
    "      else:\n",
    "          neg_scores = None\n",
    "      return pos_scores, neg_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D3eJstsehMqc"
   },
   "source": [
    "## Data loaders and train loops\n",
    "\n",
    "We aims to solve the link prediction problem, so our dataloader should iterate over the edges.\n",
    "\n",
    "Our model should aggregate the information over several layers of neighbors. This could be done usin the NeighborSampling in torch. We can define the desired number of layers and desired number of neighbors at each layer by passing fanouts list. The neighbors will be sampled uniformly.\n",
    "\n",
    "Let us use the 3-hop neighborhood with 10, 5 and 3 sampled neighbors. The general idea of such setting is that local structure plays much more role for in our network\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YcypriM3IuxL"
   },
   "outputs": [],
   "source": [
    "sampler = dgl.dataloading.MultiLayerNeighborSampler([10, 5, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NsA9om9dIvXs"
   },
   "source": [
    "The negative sampling is also required for the link prediction problem.\n",
    "We will sample one exemplar of negative edge for each positive. However, our graph is bipartite, so we need to modify standard sampler to sample only nodes of the opposite type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G2UjWKXK7Y3v"
   },
   "outputs": [],
   "source": [
    "class BipartiteUniform(dgl.dataloading.negative_sampler._BaseNegativeSampler):\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "    def _generate(self, g, eids, canonical_etype):\n",
    "        shape = dgl.backend.shape(eids)\n",
    "        dtype = dgl.backend.dtype(eids)\n",
    "        ctx = dgl.backend.context(eids)\n",
    "        shape = (shape[0] * self.k,)\n",
    "        src, _ = g.find_edges(eids, etype=canonical_etype)\n",
    "        src = dgl.backend.repeat(src, self.k, 0)\n",
    "        num_users = g.number_of_nodes() - g.ndata[\"ntype\"].sum().int().item()\n",
    "        dst_from_item = dgl.backend.randint(shape, dtype, ctx, 0, num_users)\n",
    "        dst_from_user = dgl.backend.randint(shape, dtype, ctx, 0, g.ndata[\"ntype\"].sum().int().item()) + num_users\n",
    "        dst = torch.where(g.ndata[\"ntype\"][src].view(-1) == 1, dst_from_item, dst_from_user)\n",
    "        return src, dst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vyKm9eIaJFPW"
   },
   "outputs": [],
   "source": [
    "negative_sampler = BipartiteUniform(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rs7bEL0HJFoj"
   },
   "source": [
    "Finally, one can define the dataloaders.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XIEJ6OnbmNf6"
   },
   "outputs": [],
   "source": [
    "batch_size = 512\n",
    "device = \"cpu\"\n",
    "\n",
    "train_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    train_graph,\n",
    "    train_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)\n",
    "\n",
    "val_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    val_graph,\n",
    "    val_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)\n",
    "\n",
    "test_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    test_graph,\n",
    "    test_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CGc7A0LUOhB_"
   },
   "source": [
    "Let us define the train and evaluation functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1wFpr_v0Ogis"
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import average_precision_score\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "def train(model, loader, opt, delta=0.5):\n",
    "    model.train()\n",
    "    loss_log = []\n",
    "    for (in_nodes, pos_graph, neg_graph, blocks) in tqdm(loader):\n",
    "        preds = model(pos_graph, neg_graph, blocks)\n",
    "\n",
    "        loss = (-preds[0] + preds[1] + delta).clamp(min=0).sum()\n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        loss_log.append(loss.item())\n",
    "    return loss_log\n",
    "\n",
    "\n",
    "def evaluate(model, loader, opt):\n",
    "    model.eval()\n",
    "    pred_log = []\n",
    "    gt_log = [torch.cat([torch.ones((batch_size, 1)), torch.zeros((batch_size, 1))]).data.numpy()]\n",
    "    for (in_nodes, pos_graph, neg_graph, blocks) in tqdm(loader):\n",
    "        preds = torch.cat(model(pos_graph, neg_graph, blocks))\n",
    "        pred_log.append(preds.data.cpu().numpy())\n",
    "    gt_log *= len(pred_log) - 1\n",
    "    gt_log.append(torch.cat([torch.ones((len(pred_log[-1]) // 2, 1)), torch.zeros((len(pred_log[-1]) // 2, 1))]).data.numpy())\n",
    "    return average_precision_score(np.concatenate(gt_log), np.concatenate(pred_log))\n",
    "    \n",
    "def run(model, train_loader, val_loader, test_loader, opt, n_epoch=100, delta=0.1, trainer=train):\n",
    "    ap_log = []\n",
    "    for epoch in range(n_epoch):\n",
    "        loss_log = trainer(model, train_loader, opt, delta=delta)\n",
    "        ap_log.append(evaluate(model, val_loader, opt))\n",
    "        clear_output()\n",
    "        plt.plot(ap_log)\n",
    "        plt.title(f\"Best score: {max(ap_log)} at epoch: {np.argmax(ap_log)}, last score: {ap_log[-1]} at epoch {epoch}\")\n",
    "        plt.show()\n",
    "    print(f'Test results: {evaluate(model, test_loader, opt)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 682
    },
    "id": "TR8z0K3zo1W7",
    "outputId": "5f6fefa0-3d8e-4e76-b95e-149fdc2ff551"
   },
   "outputs": [],
   "source": [
    "model = GraphSAGE(29, 8, 3, nn.ReLU(), 0.1)\n",
    "model.to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), 3e-4)\n",
    "run(model, train_loader, val_loader, test_loader, opt, 10, delta=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SbfTEa2ddOeq"
   },
   "source": [
    "We test the quality of base training. However, the rate of a movie shows how much an user like them, so it will be helpful to account such information in loss.\n",
    "\n",
    "Now we train our model with simple margin-loss. To account for the heterogeneity of links we will add low-rank positive part to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "evZofcJ1dHzu"
   },
   "outputs": [],
   "source": [
    "def train(model, loader, opt, delta=0.1):\n",
    "    model.train()\n",
    "    loss_log = []\n",
    "    for (in_nodes, pos_graph, neg_graph, blocks) in tqdm(loader):\n",
    "        preds = model(pos_graph, neg_graph, blocks)\n",
    "        loss = (-preds[0][pos_graph.edata[\"rating\"] > 3] + preds[1][pos_graph.edata[\"rating\"] > 3] + delta).clamp(min=0).sum()\n",
    "        loss += (-preds[0][pos_graph.edata[\"rating\"] <= 3] + preds[1][pos_graph.edata[\"rating\"] <= 3] + delta / 2).clamp(min=0).sum()\n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        loss_log.append(loss.item())\n",
    "    return loss_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 316
    },
    "id": "jypzPS39dHxT",
    "outputId": "f1b4e135-6b02-4afd-b3be-313ca35beec2"
   },
   "outputs": [],
   "source": [
    "model = GraphSAGE(29, 8, 3, nn.ReLU(), 0.1)\n",
    "model.to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), 3e-4)\n",
    "run(model, train_loader, val_loader, test_loader, opt, 10, delta=0.1, trainer=train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1UUj28_QxwDT"
   },
   "source": [
    "## PinSAGE sampling\n",
    "\n",
    "In this task we will define the sampler from PinSAGE paper.\n",
    "The general idea is to sample neighbors with random walks and weight the message passing with hitting count (approximation of PageRank)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IA_Pro0525oS"
   },
   "outputs": [],
   "source": [
    "class PinSAGESampler(dgl.dataloading.MultiLayerNeighborSampler):\n",
    "    def __init__(self, walk_len, num_walks_per_node, top_k, *args, **kwargs):\n",
    "        super(PinSAGESampler, self).__init__(*args, **kwargs)\n",
    "        self.walk_len = walk_len\n",
    "        self.num_walks_per_node = num_walks_per_node\n",
    "        self.top_k = top_k\n",
    "\n",
    "    def sample_frontier(self, block_id, g, seed_nodes, exclude_eids=None):\n",
    "        seed_nodes = seed_nodes[\"_N\"] if type(seed_nodes) is dict  else seed_nodes\n",
    "        walks, types = dgl.sampling.random_walk(graph, torch.LongTensor(np.repeat(seed_nodes, self.num_walks_per_node)), length=self.walk_len)\n",
    "        src = dgl.backend.reshape(walks[:, self.walk_len::self.walk_len], (-1,))\n",
    "        dst = dgl.backend.repeat(walks[:, 0], 1, 0)\n",
    "        src_mask = (src != -1)\n",
    "        src = dgl.backend.boolean_mask(src, src_mask)\n",
    "        dst = dgl.backend.boolean_mask(dst, src_mask)\n",
    "\n",
    "        neighbor_graph = dgl.graph((src, dst))\n",
    "        neighbor_graph = neighbor_graph.to_simple(return_counts=\"weight\")\n",
    "        counts = neighbor_graph.edata[\"weight\"]\n",
    "        neighbor_graph = dgl.sampling.neighbor.select_topk(neighbor_graph, self.top_k, \"weight\")\n",
    "        selected_counts = dgl.backend.gather_row(counts, neighbor_graph.edata[dgl.EID])\n",
    "        neighbor_graph.edata[\"weight\"] = selected_counts.float() / self.walk_len / self.num_walks_per_node\n",
    "        neighbor_graph.ndata[\"feat\"] = g.ndata[\"feat\"][neighbor_graph.nodes()]\n",
    "        neighbor_graph.ndata[\"ntype\"] = g.ndata[\"ntype\"][neighbor_graph.nodes()]\n",
    "        return neighbor_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y34v9xXDWbMD"
   },
   "source": [
    "Now we need to update our data loaders with new sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jLqwLnQfJBun"
   },
   "outputs": [],
   "source": [
    "sampler = PinSAGESampler(10, 10, 5, [1, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CT3LRyOa1qBL"
   },
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "device = \"cpu\"\n",
    "\n",
    "train_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    train_graph,\n",
    "    train_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)\n",
    "\n",
    "val_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    val_graph,\n",
    "    val_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)\n",
    "\n",
    "test_loader = dgl.dataloading.EdgeDataLoader(\n",
    "    test_graph,\n",
    "    test_edges,\n",
    "    sampler,\n",
    "    negative_sampler=negative_sampler,\n",
    "    device=device,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ItwvVEWIWg_7"
   },
   "source": [
    "And let us test the final model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 298
    },
    "id": "LYqxv9-RXDfO",
    "outputId": "ca31ff11-b012-4de7-a8f9-f595a236d736"
   },
   "outputs": [],
   "source": [
    "model = GraphSAGE(29, 8, 3, nn.ReLU(), 0.1)\n",
    "model.to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), 3e-4)\n",
    "run(model, train_loader, val_loader, test_loader, opt, 10, delta=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tug2ESr7dHs4"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Made 2021. Advanced GNN. Recommender Systems DGL.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}