7-4,DeepFM模型.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ca129acc",
   "metadata": {},
   "source": [
    "\n",
    "# 7-4，DeepFM网络\n",
    "\n",
    "推荐系统和广告CTR预估主流模型的演化有两条主要路线。\n",
    "\n",
    "第一条是显式建模特征交互，提升模型对交叉特征的捕获能力。(如Wide&Deep,PNN,FNN,DCN,DeepFM,AutoInt等)\n",
    "\n",
    "第二条是加入注意力机制，提升模型的自适应能力和解释性。(如DIN,DIEN,DSIN,FiBiNET,AutoInt等)\n",
    "\n",
    "在所有这些模型中，DeepFM属于性价比非常高的模型（结构简洁，计算高效，指标有竞争力）。\n",
    "\n",
    "张俊林大佬 在2019年的时候甚至建议 沿着 LR->FM->DeepFM->干点别的 这样的路线去迭代推荐系统。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4da1dc29",
   "metadata": {},
   "source": [
    "参考文档：\n",
    "\n",
    "* 《推荐系统CTR预估学习路线》：https://zhuanlan.zhihu.com/p/351078721 \n",
    "\n",
    "* criteo数据集榜单：https://paperswithcode.com/dataset/criteo\n",
    "\n",
    "* DeepFM论文： https://arxiv.org/abs/1703.04247\n",
    "\n",
    "* 《清晰易懂，基于pytorch的DeepFM的完整实验代码》： https://zhuanlan.zhihu.com/p/332786045\n",
    "\n",
    "* torch实现参考：https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/dfm.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d1ac96f",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "<font color=\"red\">\n",
    " \n",
    "公众号 **算法美食屋** 回复关键词：**pytorch**， 获取本项目源码和所用数据集百度云盘下载链接。\n",
    "    \n",
    "</font> \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b61a84a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torchkeras==3.2.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56892bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torchkeras \n",
    "print(\"torch.__version__ = \", torch.__version__)\n",
    "print(\"torchkeras.__version__ = \", torchkeras.__version__) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf724455",
   "metadata": {},
   "source": [
    "```\n",
    "torch.__version__ =  1.10.0\n",
    "torchkeras.__version__ =  3.2.3\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e254ad5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f31d986e",
   "metadata": {},
   "source": [
    "## 一，DeepFM原理解析"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2c4ee8",
   "metadata": {},
   "source": [
    "DeepFM继承了DeepWide的主体结构，将高低特征进行融合。\n",
    "\n",
    "其主要创新点有2个。\n",
    "\n",
    "一是将Wide部分替换成了 FM结构，以更有效的捕获特征交互interaction.\n",
    "\n",
    "二是FM中的隐向量 和 Deep部分的 embedding 向量共享权重，减少模型复杂性。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2917d911",
   "metadata": {},
   "source": [
    "![](https://tva1.sinaimg.cn/large/e6c9d24egy1h1brqobat6j21260nq0xl.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "350eb2f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cdcad12",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eb3bba2",
   "metadata": {},
   "source": [
    "## 二，DeepFM的pytorch实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e601ff88",
   "metadata": {},
   "source": [
    "下面是DeepFM的一个pytorch实现。\n",
    "\n",
    "除了添加了一个并行的MLP模块用于捕获隐式高阶交叉和组合特征外，其余结构基本和FM的实现完全一致。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b01f881",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from torch import nn,Tensor \n",
    "import torch.nn.functional as F \n",
    "\n",
    "class NumEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    连续特征用linear层编码\n",
    "    输入shape: [batch_size,features_num(n), d_in], # d_in 通常是1\n",
    "    输出shape: [batch_size,features_num(n), d_out]\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n: int, d_in: int, d_out: int, bias: bool = False) -> None:\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(Tensor(n, d_in, d_out))\n",
    "        self.bias = nn.Parameter(Tensor(n, d_out)) if bias else None\n",
    "        with torch.no_grad():\n",
    "            for i in range(n):\n",
    "                layer = nn.Linear(d_in, d_out)\n",
    "                self.weight[i] = layer.weight.T\n",
    "                if self.bias is not None:\n",
    "                    self.bias[i] = layer.bias\n",
    "\n",
    "    def forward(self, x_num):\n",
    "        # x_num: batch_size, features_num, d_in\n",
    "        assert x_num.ndim == 3\n",
    "        #x = x_num[..., None] * self.weight[None]\n",
    "        #x = x.sum(-2)\n",
    "        x = torch.einsum(\"bfi,fij->bfj\",x_num,self.weight)\n",
    "        if self.bias is not None:\n",
    "            x = x + self.bias[None]\n",
    "        return x\n",
    "    \n",
    "class CatEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    离散特征用Embedding层编码\n",
    "    输入shape: [batch_size,features_num], \n",
    "    输出shape: [batch_size,features_num, d_embed]\n",
    "    \"\"\"\n",
    "    def __init__(self, categories, d_embed):\n",
    "        super().__init__()\n",
    "        self.embedding = torch.nn.Embedding(sum(categories), d_embed)\n",
    "        self.offsets = nn.Parameter(\n",
    "                torch.tensor([0] + categories[:-1]).cumsum(0),requires_grad=False)\n",
    "        \n",
    "        torch.nn.init.xavier_uniform_(self.embedding.weight.data)\n",
    "\n",
    "    def forward(self, x_cat):\n",
    "        \"\"\"\n",
    "        :param x_cat: Long tensor of size ``(batch_size, features_num)``\n",
    "        \"\"\"\n",
    "        x = x_cat + self.offsets[None]\n",
    "        return self.embedding(x) \n",
    "    \n",
    "class CatLinear(nn.Module):\n",
    "    \"\"\"\n",
    "    离散特征用Embedding实现线性层（等价于先F.onehot再nn.Linear()）\n",
    "    输入shape: [batch_size,features_num], \n",
    "    输出shape: [batch_size,features_num, d_out]\n",
    "    \"\"\"\n",
    "    def __init__(self, categories, d_out=1):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Embedding(sum(categories), d_out)\n",
    "        self.bias = nn.Parameter(torch.zeros((d_out,)))\n",
    "        self.offsets = nn.Parameter(\n",
    "                torch.tensor([0] + categories[:-1]).cumsum(0),requires_grad=False)\n",
    "\n",
    "    def forward(self, x_cat):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, features_num)``\n",
    "        \"\"\"\n",
    "        x = x_cat + self.offsets[None]\n",
    "        return torch.sum(self.fc(x), dim=1) + self.bias \n",
    "    \n",
    "    \n",
    "class FMLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    FM交互项\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, reduce_sum=True):\n",
    "        super().__init__()\n",
    "        self.reduce_sum = reduce_sum\n",
    "\n",
    "    def forward(self, x): #注意：这里的x是公式中的 <v_i> * xi\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, num_features, k)``\n",
    "        \"\"\"\n",
    "        square_of_sum = torch.sum(x, dim=1) ** 2\n",
    "        sum_of_square = torch.sum(x ** 2, dim=1)\n",
    "        ix = square_of_sum - sum_of_square\n",
    "        if self.reduce_sum:\n",
    "            ix = torch.sum(ix, dim=1, keepdim=True)\n",
    "        return 0.5 * ix\n",
    "    \n",
    "    \n",
    "    \n",
    "#deep部分\n",
    "class MultiLayerPerceptron(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_layers, dropout, \n",
    "                 d_out = 1):\n",
    "        super().__init__()\n",
    "        layers = []\n",
    "        for d in d_layers:\n",
    "            layers.append(nn.Linear(d_in, d))\n",
    "            layers.append(nn.BatchNorm1d(d))\n",
    "            layers.append(nn.ReLU())\n",
    "            layers.append(nn.Dropout(p=dropout))\n",
    "            d_in = d\n",
    "        layers.append(nn.Linear(d_layers[-1], d_out))\n",
    "        self.mlp = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, d_in)``\n",
    "        \"\"\"\n",
    "        return self.mlp(x)\n",
    "    \n",
    "    \n",
    "class DeepFM(nn.Module):\n",
    "    \"\"\"\n",
    "    DeepFM模型。\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, d_numerical, categories, d_embed,\n",
    "                 deep_layers, deep_dropout,\n",
    "                 n_classes = 1):\n",
    "        \n",
    "        super().__init__()\n",
    "        if d_numerical is None:\n",
    "            d_numerical = 0\n",
    "        if categories is None:\n",
    "            categories = []\n",
    "            \n",
    "        self.categories = categories\n",
    "        self.n_classes = n_classes\n",
    "        \n",
    "        self.num_linear = nn.Linear(d_numerical,n_classes) if d_numerical else None\n",
    "        self.cat_linear = CatLinear(categories,n_classes) if categories else None\n",
    "        \n",
    "        self.num_embedding = NumEmbedding(d_numerical,1,d_embed) if d_numerical else None\n",
    "        self.cat_embedding = CatEmbedding(categories, d_embed) if categories else None\n",
    "        \n",
    "        if n_classes==1:\n",
    "            self.fm = FMLayer(reduce_sum=True)\n",
    "            self.fm_linear = None\n",
    "        else:\n",
    "            assert n_classes>=2\n",
    "            self.fm = FMLayer(reduce_sum=False)\n",
    "            self.fm_linear = nn.Linear(d_embed,n_classes)\n",
    "            \n",
    "        self.deep_in = d_numerical*d_embed+len(categories)*d_embed\n",
    "        \n",
    "        self.deep = MultiLayerPerceptron(\n",
    "            d_in= self.deep_in,\n",
    "            d_layers = deep_layers,\n",
    "            dropout = deep_dropout,\n",
    "            d_out = n_classes\n",
    "        )\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        \"\"\"\n",
    "        x_num: numerical features\n",
    "        x_cat: category features\n",
    "        \"\"\"\n",
    "        x_num,x_cat = x\n",
    "        #linear部分\n",
    "        x = 0.0\n",
    "        if self.num_linear:\n",
    "            x = x + self.num_linear(x_num) \n",
    "        if self.cat_linear:\n",
    "            x = x + self.cat_linear(x_cat)\n",
    "        \n",
    "        #fm部分\n",
    "        x_embedding = []\n",
    "        if self.num_embedding:\n",
    "            x_embedding.append(self.num_embedding(x_num[...,None]))\n",
    "        if self.cat_embedding:\n",
    "            x_embedding.append(self.cat_embedding(x_cat))\n",
    "        x_embedding = torch.cat(x_embedding,dim=1)\n",
    "        \n",
    "        if self.n_classes==1:\n",
    "            x = x + self.fm(x_embedding)\n",
    "        else: \n",
    "            x = x + self.fm_linear(self.fm(x_embedding)) \n",
    "            \n",
    "        #deep部分\n",
    "        x = x + self.deep(x_embedding.view(-1,self.deep_in))\n",
    "        \n",
    "        if self.n_classes==1:\n",
    "            x = x.squeeze(-1)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d20a06d",
   "metadata": {},
   "outputs": [],
   "source": [
    "##测试 DeepFM\n",
    "\n",
    "model = DeepFM(d_numerical = 3, categories = [4,3,2],\n",
    "        d_embed = 4, deep_layers = [20,20], deep_dropout=0.1,\n",
    "        n_classes = 1)\n",
    "x_num = torch.randn(2,3)\n",
    "x_cat = torch.randint(0,2,(2,3))\n",
    "model((x_num,x_cat)) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8316d366",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "23b6a528",
   "metadata": {},
   "source": [
    "## 三，criteo数据集完整范例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c4fce3",
   "metadata": {},
   "source": [
    "Criteo数据集是一个经典的广告点击率CTR预测数据集。\n",
    "\n",
    "这个数据集的目标是通过用户特征和广告特征来预测某条广告是否会为用户点击。\n",
    "\n",
    "数据集有13维数值特征(I1~I13)和26维类别特征(C14~C39), 共39维特征, 特征中包含着许多缺失值。\n",
    "\n",
    "训练集4000万个样本，测试集600万个样本。数据集大小超过100G.\n",
    "\n",
    "此处使用的是采样100万个样本后的cretio_small数据集。 \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa1bb01c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13245a35",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import pandas as pd \n",
    "import datetime \n",
    "\n",
    "from sklearn.model_selection import train_test_split \n",
    "\n",
    "import torch \n",
    "from torch import nn \n",
    "from torch.utils.data import Dataset,DataLoader  \n",
    "import torch.nn.functional as F \n",
    "import torchkeras \n",
    "\n",
    "def printlog(info):\n",
    "    nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "    print(\"\\n\"+\"==========\"*8 + \"%s\"%nowtime)\n",
    "    print(info+'...\\n\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e71b6b04",
   "metadata": {},
   "source": [
    "### 1，准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a89c17a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder,QuantileTransformer\n",
    "from sklearn.pipeline import Pipeline \n",
    "from sklearn.impute import SimpleImputer \n",
    "\n",
    "dfdata = pd.read_csv(\"./eat_pytorch_datasets/criteo_small.zip\",sep=\"\\t\",header=None)\n",
    "dfdata.columns = [\"label\"] + [\"I\"+str(x) for x in range(1,14)] + [\n",
    "    \"C\"+str(x) for x in range(14,40)]\n",
    "\n",
    "cat_cols = [x for x in dfdata.columns if x.startswith('C')]\n",
    "num_cols = [x for x in dfdata.columns if x.startswith('I')]\n",
    "num_pipe = Pipeline(steps = [('impute',SimpleImputer()),('quantile',QuantileTransformer())])\n",
    "\n",
    "for col in cat_cols:\n",
    "    dfdata[col]  = LabelEncoder().fit_transform(dfdata[col])\n",
    "\n",
    "dfdata[num_cols] = num_pipe.fit_transform(dfdata[num_cols])\n",
    "\n",
    "categories = [dfdata[col].max()+1 for col in cat_cols]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f30f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from torch.utils.data import Dataset,DataLoader \n",
    "\n",
    "#DataFrame转换成torch数据集Dataset, 特征分割成X_num,X_cat方式\n",
    "class DfDataset(Dataset):\n",
    "    def __init__(self,df,\n",
    "                 label_col,\n",
    "                 num_features,\n",
    "                 cat_features,\n",
    "                 categories,\n",
    "                 is_training=True):\n",
    "        \n",
    "        self.X_num = torch.tensor(df[num_features].values).float() if num_features else None\n",
    "        self.X_cat = torch.tensor(df[cat_features].values).long() if cat_features else None\n",
    "        self.Y = torch.tensor(df[label_col].values).float() \n",
    "        self.categories = categories\n",
    "        self.is_training = is_training\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.Y)\n",
    "    \n",
    "    def __getitem__(self,index):\n",
    "        if self.is_training:\n",
    "            return ((self.X_num[index],self.X_cat[index]),self.Y[index])\n",
    "        else:\n",
    "            return (self.X_num[index],self.X_cat[index])\n",
    "    \n",
    "    def get_categories(self):\n",
    "        return self.categories\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b0eba3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dftrain_val,dftest = train_test_split(dfdata,test_size=0.2)\n",
    "dftrain,dfval = train_test_split(dftrain_val,test_size=0.2)\n",
    "\n",
    "ds_train = DfDataset(dftrain,label_col = \"label\",num_features = num_cols,cat_features = cat_cols,\n",
    "                    categories = categories, is_training=True)\n",
    "\n",
    "ds_val = DfDataset(dfval,label_col = \"label\",num_features = num_cols,cat_features = cat_cols,\n",
    "                    categories = categories, is_training=True)\n",
    "\n",
    "ds_test = DfDataset(dftest,label_col = \"label\",num_features = num_cols,cat_features = cat_cols,\n",
    "                    categories = categories, is_training=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04f56e49",
   "metadata": {},
   "outputs": [],
   "source": [
    "dl_train = DataLoader(ds_train,batch_size = 2048,shuffle=True)\n",
    "dl_val = DataLoader(ds_val,batch_size = 2048,shuffle=False)\n",
    "dl_test = DataLoader(ds_test,batch_size = 2048,shuffle=False)\n",
    "\n",
    "for features,labels in dl_train:\n",
    "    break \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d67386df",
   "metadata": {},
   "source": [
    "### 2，定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d303b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_net():\n",
    "    net = DeepFM(\n",
    "        d_numerical= ds_train.X_num.shape[1],\n",
    "        categories= ds_train.get_categories(),\n",
    "        d_embed = 8, deep_layers = [128,64,32], deep_dropout=0.25,\n",
    "        n_classes = 1\n",
    "        \n",
    "    )\n",
    "    return net \n",
    "\n",
    "from torchkeras import summary\n",
    "\n",
    "net = create_net()\n",
    "summary(net,input_data=features);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be7ebf2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "dd85feec",
   "metadata": {},
   "source": [
    "```\n",
    "--------------------------------------------------------------------------\n",
    "Layer (type)                            Output Shape              Param #\n",
    "==========================================================================\n",
    "Linear-1                                     [-1, 1]                   14\n",
    "Embedding-2                              [-1, 26, 1]            1,296,709\n",
    "NumEmbedding-3                           [-1, 13, 8]                  104\n",
    "Embedding-4                              [-1, 26, 8]           10,373,672\n",
    "FMLayer-5                                    [-1, 1]                    0\n",
    "Linear-6                                   [-1, 128]               40,064\n",
    "BatchNorm1d-7                              [-1, 128]                  256\n",
    "ReLU-8                                     [-1, 128]                    0\n",
    "Dropout-9                                  [-1, 128]                    0\n",
    "Linear-10                                   [-1, 64]                8,256\n",
    "BatchNorm1d-11                              [-1, 64]                  128\n",
    "ReLU-12                                     [-1, 64]                    0\n",
    "Dropout-13                                  [-1, 64]                    0\n",
    "Linear-14                                   [-1, 32]                2,080\n",
    "BatchNorm1d-15                              [-1, 32]                   64\n",
    "ReLU-16                                     [-1, 32]                    0\n",
    "Dropout-17                                  [-1, 32]                    0\n",
    "Linear-18                                    [-1, 1]                   33\n",
    "==========================================================================\n",
    "Total params: 11,721,380\n",
    "Trainable params: 11,721,380\n",
    "Non-trainable params: 0\n",
    "--------------------------------------------------------------------------\n",
    "Input size (MB): 0.000084\n",
    "Forward/backward pass size (MB): 0.009438\n",
    "Params size (MB): 44.713516\n",
    "Estimated Total Size (MB): 44.723038\n",
    "--------------------------------------------------------------------------\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ec71ed",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "cae075b1",
   "metadata": {},
   "source": [
    "### 3，训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df4eb363",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os,sys,time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import datetime \n",
    "from tqdm import tqdm \n",
    "\n",
    "import torch\n",
    "from torch import nn \n",
    "from accelerate import Accelerator\n",
    "from copy import deepcopy\n",
    "\n",
    "\n",
    "def printlog(info):\n",
    "    nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "    print(\"\\n\"+\"==========\"*8 + \"%s\"%nowtime)\n",
    "    print(str(info)+\"\\n\")\n",
    "    \n",
    "class StepRunner:\n",
    "    def __init__(self, net, loss_fn,stage = \"train\", metrics_dict = None, \n",
    "                 optimizer = None, lr_scheduler = None,\n",
    "                 accelerator = None\n",
    "                 ):\n",
    "        self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage\n",
    "        self.optimizer,self.lr_scheduler = optimizer,lr_scheduler\n",
    "        self.accelerator = accelerator\n",
    "    \n",
    "    def __call__(self, features, labels):\n",
    "        #loss\n",
    "        preds = self.net(features)\n",
    "        loss = self.loss_fn(preds,labels)\n",
    "\n",
    "        #backward()\n",
    "        if self.optimizer is not None and self.stage==\"train\":\n",
    "            if self.accelerator is  None:\n",
    "                loss.backward()\n",
    "            else:\n",
    "                self.accelerator.backward(loss)\n",
    "            self.optimizer.step()\n",
    "            if self.lr_scheduler is not None:\n",
    "                self.lr_scheduler.step()\n",
    "            self.optimizer.zero_grad()\n",
    "            \n",
    "        #metrics\n",
    "        step_metrics = {self.stage+\"_\"+name:metric_fn(preds, labels).item() \n",
    "                        for name,metric_fn in self.metrics_dict.items()}\n",
    "        return loss.item(),step_metrics\n",
    "    \n",
    "    \n",
    "class EpochRunner:\n",
    "    def __init__(self,steprunner):\n",
    "        self.steprunner = steprunner\n",
    "        self.stage = steprunner.stage\n",
    "        self.steprunner.net.train() if self.stage==\"train\" else self.steprunner.net.eval()\n",
    "        \n",
    "    def __call__(self,dataloader):\n",
    "        total_loss,step = 0,0\n",
    "        loop = tqdm(enumerate(dataloader), total =len(dataloader))\n",
    "        for i, batch in loop:\n",
    "            features,labels = batch\n",
    "            if self.stage==\"train\":\n",
    "                loss, step_metrics = self.steprunner(features,labels)\n",
    "            else:\n",
    "                with torch.no_grad():\n",
    "                    loss, step_metrics = self.steprunner(features,labels)\n",
    "\n",
    "            step_log = dict({self.stage+\"_loss\":loss},**step_metrics)\n",
    "\n",
    "            total_loss += loss\n",
    "            step+=1\n",
    "            if i!=len(dataloader)-1:\n",
    "                loop.set_postfix(**step_log)\n",
    "            else:\n",
    "                epoch_loss = total_loss/step\n",
    "                epoch_metrics = {self.stage+\"_\"+name:metric_fn.compute().item() \n",
    "                                 for name,metric_fn in self.steprunner.metrics_dict.items()}\n",
    "                epoch_log = dict({self.stage+\"_loss\":epoch_loss},**epoch_metrics)\n",
    "                loop.set_postfix(**epoch_log)\n",
    "\n",
    "                for name,metric_fn in self.steprunner.metrics_dict.items():\n",
    "                    metric_fn.reset()\n",
    "        return epoch_log\n",
    "\n",
    "class KerasModel(torch.nn.Module):\n",
    "    def __init__(self,net,loss_fn,metrics_dict=None,optimizer=None,lr_scheduler = None):\n",
    "        super().__init__()\n",
    "        self.accelerator = Accelerator()\n",
    "        self.history = {}\n",
    "        \n",
    "        self.net = net\n",
    "        self.loss_fn = loss_fn\n",
    "        self.metrics_dict = nn.ModuleDict(metrics_dict) \n",
    "        \n",
    "        self.optimizer = optimizer if optimizer is not None else torch.optim.Adam(\n",
    "            self.parameters(), lr=1e-2)\n",
    "        self.lr_scheduler = lr_scheduler\n",
    "        \n",
    "        self.net,self.loss_fn,self.metrics_dict,self.optimizer = self.accelerator.prepare(\n",
    "            self.net,self.loss_fn,self.metrics_dict,self.optimizer)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.net:\n",
    "            return self.net.forward(x)\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "\n",
    "    def fit(self, train_data, val_data=None, epochs=10, ckpt_path='checkpoint.pt', \n",
    "            patience=5, monitor=\"val_loss\", mode=\"min\"):\n",
    "        \n",
    "        train_data = self.accelerator.prepare(train_data)\n",
    "        val_data = self.accelerator.prepare(val_data) if val_data else []\n",
    "\n",
    "        for epoch in range(1, epochs+1):\n",
    "            printlog(\"Epoch {0} / {1}\".format(epoch, epochs))\n",
    "            \n",
    "            # 1，train -------------------------------------------------  \n",
    "            train_step_runner = StepRunner(net = self.net,stage=\"train\",\n",
    "                    loss_fn = self.loss_fn,metrics_dict=deepcopy(self.metrics_dict),\n",
    "                    optimizer = self.optimizer, lr_scheduler = self.lr_scheduler,\n",
    "                    accelerator = self.accelerator)\n",
    "            train_epoch_runner = EpochRunner(train_step_runner)\n",
    "            train_metrics = train_epoch_runner(train_data)\n",
    "            \n",
    "            for name, metric in train_metrics.items():\n",
    "                self.history[name] = self.history.get(name, []) + [metric]\n",
    "\n",
    "            # 2，validate -------------------------------------------------\n",
    "            if val_data:\n",
    "                val_step_runner = StepRunner(net = self.net,stage=\"val\",\n",
    "                    loss_fn = self.loss_fn,metrics_dict=deepcopy(self.metrics_dict),\n",
    "                    accelerator = self.accelerator)\n",
    "                val_epoch_runner = EpochRunner(val_step_runner)\n",
    "                with torch.no_grad():\n",
    "                    val_metrics = val_epoch_runner(val_data)\n",
    "                val_metrics[\"epoch\"] = epoch\n",
    "                for name, metric in val_metrics.items():\n",
    "                    self.history[name] = self.history.get(name, []) + [metric]\n",
    "            \n",
    "            # 3，early-stopping -------------------------------------------------\n",
    "            arr_scores = self.history[monitor]\n",
    "            best_score_idx = np.argmax(arr_scores) if mode==\"max\" else np.argmin(arr_scores)\n",
    "            if best_score_idx==len(arr_scores)-1:\n",
    "                torch.save(self.net.state_dict(),ckpt_path)\n",
    "                print(\"<<<<<< reach best {0} : {1} >>>>>>\".format(monitor,\n",
    "                     arr_scores[best_score_idx]),file=sys.stderr)\n",
    "            if len(arr_scores)-best_score_idx>patience:\n",
    "                print(\"<<<<<< {} without improvement in {} epoch, early stopping >>>>>>\".format(\n",
    "                    monitor,patience),file=sys.stderr)\n",
    "                break \n",
    "                \n",
    "        self.net.load_state_dict(torch.load(ckpt_path))\n",
    "            \n",
    "        return pd.DataFrame(self.history)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def evaluate(self, val_data):\n",
    "        val_data = self.accelerator.prepare(val_data)\n",
    "        val_step_runner = StepRunner(net = self.net,stage=\"val\",\n",
    "                    loss_fn = self.loss_fn,metrics_dict=deepcopy(self.metrics_dict),\n",
    "                    accelerator = self.accelerator)\n",
    "        val_epoch_runner = EpochRunner(val_step_runner)\n",
    "        val_metrics = val_epoch_runner(val_data)\n",
    "        return val_metrics\n",
    "        \n",
    "       \n",
    "    @torch.no_grad()\n",
    "    def predict(self, dataloader):\n",
    "        dataloader = self.accelerator.prepare(dataloader)\n",
    "        result = torch.cat([self.forward(t[0]) for t in dataloader])\n",
    "        return result.data\n",
    "              "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f5987a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchkeras.metrics import AUC\n",
    "\n",
    "loss_fn = nn.BCEWithLogitsLoss()\n",
    "\n",
    "metrics_dict = {\"auc\":AUC()}\n",
    "\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=0.002, weight_decay=0.001) \n",
    "\n",
    "model = KerasModel(net,\n",
    "                   loss_fn = loss_fn,\n",
    "                   metrics_dict= metrics_dict,\n",
    "                   optimizer = optimizer\n",
    "                  )      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3ec3f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfhistory = model.fit(train_data=dl_train,val_data=dl_val,epochs=50, patience=5,\n",
    "                      monitor = \"val_auc\",mode=\"max\",ckpt_path='checkpoint.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26087443",
   "metadata": {},
   "source": [
    "### 4，评估模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8937aa90",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'svg'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def plot_metric(dfhistory, metric):\n",
    "    train_metrics = dfhistory[\"train_\"+metric]\n",
    "    val_metrics = dfhistory['val_'+metric]\n",
    "    epochs = range(1, len(train_metrics) + 1)\n",
    "    plt.plot(epochs, train_metrics, 'bo--')\n",
    "    plt.plot(epochs, val_metrics, 'ro-')\n",
    "    plt.title('Training and validation '+ metric)\n",
    "    plt.xlabel(\"Epochs\")\n",
    "    plt.ylabel(metric)\n",
    "    plt.legend([\"train_\"+metric, 'val_'+metric])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17f0ebe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(dfhistory,\"loss\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5adcb4b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(dfhistory,\"auc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27826a5a",
   "metadata": {},
   "source": [
    "### 5，使用模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1fd052c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "preds = torch.sigmoid(model.predict(dl_val))\n",
    "labels = torch.cat([x[-1] for x in dl_val])\n",
    "\n",
    "val_auc = roc_auc_score(labels.cpu().numpy(),preds.cpu().numpy())\n",
    "print(val_auc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df8aa18e",
   "metadata": {},
   "source": [
    "### 6，保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c578b927",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.net.state_dict(),\"best_deepfm.pt\")\n",
    "net_clone = create_net()\n",
    "net_clone.load_state_dict(torch.load(\"best_deepfm.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2948189d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "net_clone.eval()\n",
    "preds = torch.cat([torch.sigmoid(net_clone(x[0])).data for x in dl_val]) \n",
    "labels = torch.cat([x[-1] for x in dl_val])\n",
    "\n",
    "val_auc = roc_auc_score(labels.cpu().numpy(),preds.cpu().numpy())\n",
    "print(val_auc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ce2aa6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0d010313",
   "metadata": {},
   "source": [
    "**如果本书对你有所帮助，想鼓励一下作者，记得给本项目加一颗星星star⭐️，并分享给你的朋友们喔😊!** \n",
    "\n",
    "如果对本书内容理解上有需要进一步和作者交流的地方，欢迎在公众号\"算法美食屋\"下留言。作者时间和精力有限，会酌情予以回复。\n",
    "\n",
    "也可以在公众号后台回复关键字：**加群**，加入读者交流群和大家讨论。\n",
    "\n",
    "![算法美食屋logo.png](https://tva1.sinaimg.cn/large/e6c9d24egy1h41m2zugguj20k00b9q46.jpg)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb,md",
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}