collaborative-filtering.ipynb

{
 "metadata": {
  "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.8.5-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python_defaultSpec_1601666261819",
   "display_name": "Python 3.8.5 64-bit"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Collaborative filtering\n",
    "\n",
    "An attempt to match fastai's performance on [MovieLens 100K](https://grouplens.org/datasets/movielens/100k/), but with pure pytorch, not using fastai's library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "import fastai\n",
    "import fastai.collab\n",
    "import fastai.datasets\n",
    "import fastai.tabular.transform\n",
    "import math\n",
    "import numpy\n",
    "import os\n",
    "import pandas\n",
    "import time\n",
    "import torch\n",
    "from torch import nn\n",
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Seems to be much faster than GPU for this application. I think the model is simple enough\n",
    "# that the overhead of transferring data between CPU and GPU RAM dominates any speedup.\n",
    "dev = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import zipfile\n",
    "zip_path = fastai.datasets.download_data(\"http://files.grouplens.org/datasets/movielens/ml-100k.zip\", ext=\"\")\n",
    "dest_dir = zip_path.parent\n",
    "data_dir = os.path.splitext(zip_path)[0]\n",
    "if not os.path.exists(data_dir):\n",
    "    with zipfile.ZipFile(zip_path) as zf:\n",
    "        zf.extractall(dest_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "n_item: 1682, n_user: 943\n"
    }
   ],
   "source": [
    "col_names = (\"user\", \"item\", \"rating\", \"timestamp\")\n",
    "train_df = pandas.read_csv(os.path.join(data_dir, \"ua.base\"), sep=\"\\t\", names=col_names)\n",
    "test_df = pandas.read_csv(os.path.join(data_dir, \"ua.test\"), sep=\"\\t\", names=col_names)\n",
    "concat_df = pandas.concat((train_df, test_df))\n",
    "n_item = concat_df[\"item\"].max()\n",
    "n_user = concat_df[\"user\"].max()\n",
    "print(f\"n_item: {n_item}, n_user: {n_user}\")\n",
    "\n",
    "class MovieLensDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, df: pandas.DataFrame, device: torch.device):\n",
    "        # Indices into embeddings need to have dtype \"long\".\n",
    "        self.ids_tensor = torch.tensor(df[[\"user\", \"item\"]].to_numpy(), dtype=torch.long, device=device)\n",
    "        # This .squeeze() is important! Passing things with mismatched shapes to MSELoss results in surprising behavior\n",
    "        # and poor fitting.\n",
    "        self.ratings_tensor = torch.tensor(df[[\"rating\"]].to_numpy(), dtype=torch.float, device=device).squeeze()\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.ids_tensor)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.ids_tensor[idx], self.ratings_tensor[idx]\n",
    "        \n",
    "\n",
    "train_dataset = MovieLensDataset(train_df, dev)\n",
    "test_dataset = MovieLensDataset(test_df, dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 64 # Copied from fastai\n",
    "num_epochs = 5  # Copied from fastai\n",
    "lr = 1e-2       # Set from tuning by hand for my model\n",
    "wd = 1e-1       # Copied from fastai\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "test_inputs, test_targets = test_dataset[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## fastai benchmark\n",
    "fastai claims state of the art performance, so let's start with that and see how it does.\n",
    "This is based on <https://github.com/fastai/course-v3/blob/master/nbs/dl1/lesson4-collab.ipynb>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": "<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th>user</th>\n      <th>item</th>\n      <th>target</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <td>761</td>\n      <td>840</td>\n      <td>4.0</td>\n    </tr>\n    <tr>\n      <td>393</td>\n      <td>727</td>\n      <td>3.0</td>\n    </tr>\n    <tr>\n      <td>318</td>\n      <td>384</td>\n      <td>3.0</td>\n    </tr>\n    <tr>\n      <td>629</td>\n      <td>392</td>\n      <td>4.0</td>\n    </tr>\n    <tr>\n      <td>190</td>\n      <td>24</td>\n      <td>3.0</td>\n    </tr>\n  </tbody>\n</table>"
     },
     "metadata": {}
    }
   ],
   "source": [
    "# Hacked up copy of CollabDataBunch.from_df because I want to use test_df as the validation set.\n",
    "user_name   = concat_df.columns[0]\n",
    "item_name   = concat_df.columns[1]\n",
    "rating_name = concat_df.columns[2]\n",
    "cat_names = [user_name,item_name]\n",
    "num_train = len(train_df)\n",
    "src = (fastai.collab.CollabList.from_df(concat_df, cat_names=cat_names, procs=fastai.tabular.transform.Categorify)\n",
    "        .split_by_idxs(train_idx=numpy.arange(num_train), valid_idx=numpy.arange(num_train, num_train + len(test_df)))\n",
    "        .label_from_df(cols=rating_name))\n",
    "data_bunch = src.databunch(path=\".\", bs=batch_size, val_bs=batch_size, device=dev)\n",
    "assert len(data_bunch.dl(fastai.basic_data.DatasetType.Train).dl.dataset.x) == num_train\n",
    "data_bunch.show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": "<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: left;\">\n      <th>epoch</th>\n      <th>train_loss</th>\n      <th>valid_loss</th>\n      <th>time</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <td>0</td>\n      <td>0.923776</td>\n      <td>1.005359</td>\n      <td>00:07</td>\n    </tr>\n    <tr>\n      <td>1</td>\n      <td>0.850828</td>\n      <td>0.957647</td>\n      <td>00:07</td>\n    </tr>\n    <tr>\n      <td>2</td>\n      <td>0.795858</td>\n      <td>0.906226</td>\n      <td>00:07</td>\n    </tr>\n    <tr>\n      <td>3</td>\n      <td>0.609235</td>\n      <td>0.881616</td>\n      <td>00:07</td>\n    </tr>\n    <tr>\n      <td>4</td>\n      <td>0.486054</td>\n      <td>0.880213</td>\n      <td>00:08</td>\n    </tr>\n  </tbody>\n</table>",
      "text/plain": "<IPython.core.display.HTML object>"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fastai_learn = fastai.collab.collab_learner(data_bunch, n_factors=40, y_range=[0.5, 5.5], wd=wd)\n",
    "fastai_learn.fit_one_cycle(num_epochs, lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "<IPython.core.display.HTML object>",
      "text/html": "\n    <div>\n        <style>\n            /* Turns off some styling */\n            progress {\n                /* gets rid of default border in Firefox and Opera. */\n                border: none;\n                /* Needs to be in here for Safari polyfill so background images work as expected. */\n                background-size: auto;\n            }\n            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n                background: #F44336;\n            }\n        </style>\n      <progress value='0' class='' max='148' style='width:300px; height:20px; vertical-align: middle;'></progress>\n      \n    </div>\n    "
     },
     "metadata": {}
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.8802131414413452"
     },
     "metadata": {},
     "execution_count": 33
    }
   ],
   "source": [
    "fastai_pred = fastai_learn.get_preds(ds_type=fastai.data_block.DatasetType.Valid)\n",
    "nn.functional.mse_loss(*fastai_pred).item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## My own implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Fitter:\n",
    "    def __init__(self, model: nn.Module, loss_func: nn.Module, optim: torch.optim.Optimizer):\n",
    "        self.model = model\n",
    "        self.loss_func = loss_func\n",
    "        self.optim = optim\n",
    "        self.losses = []\n",
    "\n",
    "    def fit(self, num_epochs: int):\n",
    "        print(\"epoch | train_loss | test_loss | time\")\n",
    "        for epoch in range(num_epochs):\n",
    "            start = time.time()\n",
    "            train_loss = torch.tensor([0.0], dtype=float, device=dev)\n",
    "            self.model.train()\n",
    "            for batch_idx, (inputs, targets) in enumerate(train_loader, 0):\n",
    "                train_loss += self._one_batch(inputs, targets)\n",
    "\n",
    "            self.model.eval()\n",
    "            with torch.no_grad():\n",
    "                # Convert from 1-based to 0-based index.\n",
    "                pred = self.model(test_inputs[:, 0] - 1, test_inputs[:, 1] - 1)\n",
    "                test_loss = self.loss_func(pred, test_targets).item()\n",
    "\n",
    "            num_batches = batch_idx + 1\n",
    "            self.losses.append((train_loss / num_batches, test_loss))\n",
    "            print(\"%5d |      %.3f |     %.3f |   %ds |\" % (\n",
    "                epoch,\n",
    "                self.losses[-1][0],\n",
    "                self.losses[-1][1],\n",
    "                int(time.time() - start)))\n",
    "\n",
    "    def _one_batch(self, inputs: torch.tensor, targets: torch.tensor) -> torch.tensor:\n",
    "        self.optim.zero_grad()\n",
    "        # Convert from 1-based to 0-based index.\n",
    "        users, items = inputs[:, 0] - 1, inputs[:, 1] - 1\n",
    "        pred = self.model(users, items)\n",
    "        loss = self.loss_func(pred, targets)\n",
    "        loss.backward()\n",
    "        self.optim.step()\n",
    "        return loss.detach()\n",
    "\n",
    "class FitterOneCycle(Fitter):\n",
    "    def __init__(self, *args):\n",
    "        super().__init__(*args)\n",
    "        self.scheduler = None\n",
    "    \n",
    "    def fit(self, num_epochs: int):\n",
    "        self.scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "            self.optim, max_lr=self.optim.defaults[\"lr\"],\n",
    "            final_div_factor=25e4, # fastai uses 25e4\n",
    "            epochs=num_epochs,\n",
    "            steps_per_epoch=math.ceil(len(train_df) / batch_size))\n",
    "        super().fit(num_epochs)\n",
    "\n",
    "    def _one_batch(self, inputs: torch.tensor, targets: torch.tensor) -> torch.tensor:\n",
    "        loss = super()._one_batch(inputs, targets)\n",
    "        self.scheduler.step()\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trunc_normal_(x: torch.tensor, mean: float=0., std: float=1.) -> torch.tensor:\n",
    "    # From https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/12\n",
    "    return x.normal_().fmod_(2).mul_(std).add_(mean)\n",
    "\n",
    "class ScaledDotProdBias(nn.Module):\n",
    "    \"\"\"Same as DotProdBias, but scale the output to be within y_range.\"\"\"\n",
    "    def __init__(self, n_user: int, n_item: int, embedding_dim: int, y_range: typing.Tuple[int, int], trunc_normal: bool=False):\n",
    "        super().__init__()\n",
    "        self.user_emb = nn.Embedding(num_embeddings=n_user, embedding_dim=embedding_dim)\n",
    "        self.item_emb = nn.Embedding(num_embeddings=n_item, embedding_dim=embedding_dim)\n",
    "        self.user_bias = nn.Embedding(num_embeddings=n_user, embedding_dim=1)\n",
    "        self.item_bias = nn.Embedding(num_embeddings=n_item, embedding_dim=1)\n",
    "        if trunc_normal:\n",
    "            # Based on\n",
    "            # https://github.com/fastai/fastai1/blob/6a5102ef7bdefa9058d0481ab311f48b21cbc6fc/fastai/layers.py#L285\n",
    "            for e in (self.user_emb, self.item_emb, self.user_bias, self.item_bias):\n",
    "                with torch.no_grad(): trunc_normal_(e.weight, std=0.01)\n",
    "        self.y_min, self.y_max = y_range\n",
    "    \n",
    "    def forward(self, users: torch.LongTensor, items: torch.LongTensor) -> torch.FloatTensor:\n",
    "        dot_prods = (self.user_emb(users) * self.item_emb(items)).sum(dim=1)\n",
    "        # Squeezing here is important! The embedding returns tensor of shape [batch_size, 1].\n",
    "        # I initially forgot this and was getting mediocre performance.\n",
    "        u_bias = self.user_bias(users).squeeze(dim=1)\n",
    "        i_bias = self.item_bias(items).squeeze(dim=1)\n",
    "        biased = dot_prods + u_bias + i_bias\n",
    "        return self.y_min + (self.y_max - self.y_min) * nn.functional.sigmoid(biased)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "My first attempt at a model didn't have the trunc_normal initialization, as it's a bit hidden in the fastai\n",
    "library.\n",
    "\n",
    "Fastai references [An Exploration of Word Embedding Initialization in Deep-Learning Tasks](https://arxiv.org/abs/1711.09160) as justification for the trunc_normal initialization. It seems the main benefit is to limit the variance since \"High-variance initialization prevents the network to use the space of embeddings and forces it to use other free parameters to accomplish the task.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "epoch | train_loss | test_loss | time\n    0 |      5.725 |     4.074 |   2s |\n    1 |      1.519 |     0.993 |   2s |\n    2 |      0.789 |     0.916 |   2s |\n    3 |      0.637 |     0.885 |   2s |\n    4 |      0.489 |     0.887 |   2s |\n"
    }
   ],
   "source": [
    "model = ScaledDotProdBias(n_user, n_item, 40, (0.5, 5.5)).to(dev)\n",
    "fitter_one_cycle = FitterOneCycle(\n",
    "    model,\n",
    "    nn.MSELoss(),\n",
    "    # betas copied from fastai, but I don't think it really made a difference vs\n",
    "    # torch's defaults.\n",
    "    torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.99), weight_decay=wd))\n",
    "fitter_one_cycle.fit(num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see how it compres to trunc_normal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "trunc_normal=True\nepoch | train_loss | test_loss | time\n    0 |      1.161 |     1.005 |   2s |\n    1 |      0.889 |     0.950 |   2s |\n    2 |      0.780 |     0.912 |   2s |\n    3 |      0.631 |     0.881 |   2s |\n    4 |      0.483 |     0.883 |   2s |\n"
    }
   ],
   "source": [
    "print(\"trunc_normal=True\")\n",
    "model = ScaledDotProdBias(n_user, n_item, 40, (0.5, 5.5), trunc_normal=True).to(dev)\n",
    "fitter_one_cycle = FitterOneCycle(\n",
    "    model,\n",
    "    nn.MSELoss(),\n",
    "    torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.99), weight_decay=wd))\n",
    "fitter_one_cycle.fit(num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initialization really helps in epoch 0, and after that the difference is small but persists.\n",
    "\n",
    "After all this, test_loss is basically the same as fastai's.\n",
    "Sometimes a bit higher or lower if I repeat the training of both, but always within 0.01.\n",
    "So, success!"
   ]
  }
 ]
}