{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.dcn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DCN\n",
    "> A pytorch implementation of Deep & Cross Network."
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "## v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "\n",
    "from recohut.models.layers.common import FeaturesEmbedding, MultiLayerPerceptron\n",
    "\n",
    "\n",
    "class CrossNetwork(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, num_layers):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.w = torch.nn.ModuleList([\n",
    "            torch.nn.Linear(input_dim, 1, bias=False) for _ in range(num_layers)\n",
    "        ])\n",
    "        self.b = torch.nn.ParameterList([\n",
    "            torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``\n",
    "        \"\"\"\n",
    "        x0 = x\n",
    "        for i in range(self.num_layers):\n",
    "            xw = self.w[i](x)\n",
    "            x = x0 * xw + self.b[i] + x\n",
    "        return x\n",
    "\n",
    "class DCN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Deep & Cross Network.\n",
    "    Reference:\n",
    "        R Wang, et al. Deep & Cross Network for Ad Click Predictions, 2017.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, num_layers, mlp_dims, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)\n",
    "        self.embed_output_dim = len(field_dims) * embed_dim\n",
    "        self.cn = CrossNetwork(self.embed_output_dim, num_layers)\n",
    "        self.mlp = MultiLayerPerceptron(self.embed_output_dim, mlp_dims, dropout, output_layer=False)\n",
    "        self.linear = torch.nn.Linear(mlp_dims[-1] + self.embed_output_dim, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        embed_x = self.embedding(x).view(-1, self.embed_output_dim)\n",
    "        x_l1 = self.cn(embed_x)\n",
    "        h_l2 = self.mlp(embed_x)\n",
    "        x_stack = torch.cat([x_l1, h_l2], dim=1)\n",
    "        p = self.linear(x_stack)\n",
    "        return torch.sigmoid(p.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "class DNN(nn.Module):\n",
    "    def __init__(self, inputs_dim, hidden_units, dropout_rate):\n",
    "        super(DNN, self).__init__()\n",
    "        self.inputs_dim = inputs_dim\n",
    "        self.hidden_units = hidden_units\n",
    "        self.dropout = nn.Dropout(dropout_rate)\n",
    "\n",
    "        self.hidden_units = [inputs_dim] + list(self.hidden_units)\n",
    "        self.linear = nn.ModuleList([\n",
    "            nn.Linear(self.hidden_units[i], self.hidden_units[i+1]) for i in range(len(self.hidden_units)-1)\n",
    "        ])\n",
    "        for name, tensor in self.linear.named_parameters():\n",
    "            if 'weight' in name:\n",
    "                nn.init.normal_(tensor, mean=0, std=0.0001)\n",
    "\n",
    "        # self.bn = nn.ModuleList([\n",
    "        #     nn.Linear(self.hidden_units[i], self.hidden_units[i + 1]) for i in range(len(self.hidden_units) - 1)\n",
    "        # ])\n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, X):\n",
    "        inputs = X\n",
    "        for i in range(len(self.linear)):\n",
    "            fc = self.linear[i](inputs)\n",
    "            fc = self.activation(fc)\n",
    "            fc - self.dropout(fc)\n",
    "            inputs = fc\n",
    "        return inputs\n",
    "\n",
    "\n",
    "class CrossNet(nn.Module):\n",
    "    def __init__(self, in_features, layer_num=2, parameterization='vector', seed=2022):\n",
    "        super(CrossNet, self).__init__()\n",
    "        self.layer_num = layer_num\n",
    "        self.parameterization = parameterization\n",
    "        if self.parameterization == 'vector':\n",
    "            self.kernels = nn.Parameter(torch.Tensor(self.layer_num, in_features, 1))\n",
    "        elif self.parameterization == 'matrix':\n",
    "            self.kernels = nn.Parameter(torch.Tensor(self.layer_num, in_features, in_features))\n",
    "        self.bias = nn.Parameter(torch.Tensor(self.layer_num, in_features, 1))\n",
    "\n",
    "        for i in range(self.kernels.shape[0]):\n",
    "            nn.init.xavier_normal_(self.kernels[i])\n",
    "        for i in range(self.bias.shape[0]):\n",
    "            nn.init.zeros_(self.bias[0])\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        x_0 = inputs.unsqueeze(2)\n",
    "        x_1 = x_0\n",
    "        for i in range(self.layer_num):\n",
    "            if self.parameterization == 'vector':\n",
    "                x1_w = torch.tensordot(x_1, self.kernels[i], dims=([1], [0]))\n",
    "                dot_ = torch.matmul(x_0, x1_w)\n",
    "                x_1 = dot_ + self.bias[i] + x_1\n",
    "            else:\n",
    "                x1_w = torch.tensordot(self.kernels[i], x_1)\n",
    "                dot_ = x1_w + self.bias[i]\n",
    "                x_1 = x_0 * dot_ + x_1\n",
    "        x_1 = torch.squeeze(x_1, dim=2)\n",
    "        return x_1\n",
    "\n",
    "\n",
    "class DCNv2(nn.Module):\n",
    "    \"\"\"DCN model implementation in pytorch\n",
    "\n",
    "    Reference:\n",
    "        1. https://github.com/huangjunheng/recommendation_model/blob/master/DCN/dcn.py\n",
    "    \"\"\"\n",
    "    def __init__(self, feat_size, embedding_size, linear_feature_columns, dnn_feature_columns, cross_num=2,\n",
    "                 cross_param='vector', dnn_hidden_units=(128, 128,), init_std=0.0001, seed=2022, l2_reg=0.00001,\n",
    "                 drop_rate=0.5):\n",
    "        super(DCNv2, self).__init__()\n",
    "        self.feat_size = feat_size\n",
    "        self.embedding_size = embedding_size\n",
    "        self.dnn_hidden_units = dnn_hidden_units\n",
    "        self.cross_num = 2\n",
    "        self.cross_param = cross_param\n",
    "        self.drop_rate = drop_rate\n",
    "        self.l2_reg = 0.00001\n",
    "\n",
    "        self.act = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(drop_rate)\n",
    "\n",
    "        self.dense_feature_columns = list(filter(lambda x:x[1]=='dense', dnn_feature_columns))\n",
    "        self.sparse_feature_columns = list(filter(lambda x:x[1]=='sparse', dnn_feature_columns))\n",
    "\n",
    "        self.embedding_dic = nn.ModuleDict({feat[0]:nn.Embedding(feat_size[feat[0]], self.embedding_size, sparse=False)\n",
    "                                            for feat in self.sparse_feature_columns})\n",
    "\n",
    "        self.feature_index = defaultdict(int)\n",
    "        start = 0\n",
    "        for feat in self.feat_size:\n",
    "            self.feature_index[feat] = start\n",
    "            start += 1\n",
    "\n",
    "        inputs_dim = len(self.dense_feature_columns)+self.embedding_size*len(self.sparse_feature_columns)\n",
    "        self.dnn = DNN(inputs_dim,self.dnn_hidden_units, 0.5)\n",
    "        self.crossnet = CrossNet(inputs_dim, layer_num=self.cross_num, parameterization=self.cross_param)\n",
    "        self.dnn_linear = nn.Linear(inputs_dim+dnn_hidden_units[-1], 1, bias=False)\n",
    "        dnn_hidden_units = [len(feat_size)] + list(dnn_hidden_units) + [1]\n",
    "        self.linear = nn.ModuleList([\n",
    "            nn.Linear(dnn_hidden_units[i], dnn_hidden_units[i+1]) for i in range(len(dnn_hidden_units)-1)\n",
    "        ])\n",
    "        for name, tensor in self.linear.named_parameters():\n",
    "            if 'weight' in name:\n",
    "                nn.init.normal_(tensor, mean=0, std=init_std)\n",
    "\n",
    "    def forward(self, X):\n",
    "        logit = X\n",
    "        for i in range(len(self.linear)):\n",
    "            fc = self.linear[i](logit)\n",
    "            fc = self.act(fc)\n",
    "            fc = self.dropout(fc)\n",
    "            logit = fc\n",
    "\n",
    "        sparse_embedding = [self.embedding_dic[feat[0]](X[:, self.feature_index[feat[0]]].long()).reshape(X.shape[0], 1, -1)\n",
    "                            for feat in self.sparse_feature_columns]\n",
    "        dense_values = [X[:, self.feature_index[feat[0]]].reshape(-1, 1) for feat in self.dense_feature_columns]\n",
    "        dense_input = torch.cat(dense_values, dim=1)\n",
    "        sparse_input = torch.cat(sparse_embedding, dim=1)\n",
    "        sparse_input = torch.flatten(sparse_input, start_dim=1)\n",
    "        dnn_input = torch.cat((dense_input, sparse_input), dim=1)\n",
    "\n",
    "        # print('sparse input size', sparse_input.shape)\n",
    "        # print('dense input size', dense_input.shape)\n",
    "        # print('dnn input size', dnn_input.shape)\n",
    "\n",
    "        deep_out = self.dnn(dnn_input)\n",
    "        cross_out = self.crossnet(dnn_input)\n",
    "        stack_out = torch.cat((cross_out, deep_out), dim=-1)\n",
    "\n",
    "        logit += self.dnn_linear(stack_out)\n",
    "        #print('logit size', logit.shape)\n",
    "        y_pred = torch.sigmoid(logit)\n",
    "        #print('y_pred', y_pred.shape)\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- R Wang, et al. Deep & Cross Network for Ad Click Predictions, 2017.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/dcn.py\n",
    "- https://github.com/huangjunheng/recommendation_model/blob/master/DCN/dcn.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p recohut"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
