{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.layers.common"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Common Layers\n",
    "> Common layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FeaturesLinear(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    References:\n",
    "        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py\n",
    "    \"\"\"\n",
    "    def __init__(self, field_dims, output_dim=1):\n",
    "        super().__init__()\n",
    "        self.fc = torch.nn.Embedding(sum(field_dims), output_dim)\n",
    "        self.bias = torch.nn.Parameter(torch.zeros((output_dim,)))\n",
    "        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        x = x + x.new_tensor(self.offsets).unsqueeze(0)\n",
    "        return torch.sum(self.fc(x), dim=1) + self.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FeaturesEmbedding(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    References:\n",
    "        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py\n",
    "    \"\"\"\n",
    "    def __init__(self, field_dims, embed_dim):\n",
    "        super().__init__()\n",
    "        self.embedding = torch.nn.Embedding(sum(field_dims), embed_dim)\n",
    "        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)\n",
    "        torch.nn.init.xavier_uniform_(self.embedding.weight.data)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        x = x + x.new_tensor(self.offsets).unsqueeze(0)\n",
    "        return self.embedding(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class MultiLayerPerceptron(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    References:\n",
    "        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, embed_dims, dropout, output_layer=True):\n",
    "        super().__init__()\n",
    "        layers = list()\n",
    "        for embed_dim in embed_dims:\n",
    "            layers.append(torch.nn.Linear(input_dim, embed_dim))\n",
    "            layers.append(torch.nn.BatchNorm1d(embed_dim))\n",
    "            layers.append(torch.nn.ReLU())\n",
    "            layers.append(torch.nn.Dropout(p=dropout))\n",
    "            input_dim = embed_dim\n",
    "        if output_layer:\n",
    "            layers.append(torch.nn.Linear(input_dim, 1))\n",
    "        self.mlp = torch.nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, embed_dim)``\n",
    "        \"\"\"\n",
    "        return self.mlp(x)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
