120分钟吃掉DIEN深度兴趣演化网络.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a93a7bc",
   "metadata": {},
   "source": [
    "# DIEN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8e56579",
   "metadata": {},
   "source": [
    "阿里妈妈在CTR预估领域有3篇比较有名的文章。\n",
    "\n",
    "2017年的深度兴趣网络, DIN(DeepInterestNetwork)。 \n",
    "\n",
    "2018年的深度兴趣演化网络, DIEN(DeepInterestEvolutionNetWork)。\n",
    "\n",
    "2019年的深度会话兴趣网络, DSIN(DeepSessionInterestNetWork)。\n",
    "\n",
    "这3篇文章的主要思想和相互关系用一句话分别概括如下：\n",
    "\n",
    "第1篇DIN说，用户的行为日志中只有一部分和当前候选广告有关。可以利用Attention机制从用户行为日志中建模出和当前候选广告相关的用户兴趣表示。我们试过涨点了嘻嘻嘻。\n",
    "\n",
    "第2篇DIEN说，用户最近的行为可能比较远的行为更加重要。可以用循环神经网络GRU建模用户兴趣随时间的演化。我们试过也涨点了嘿嘿嘿。\n",
    "\n",
    "第3篇DSIN说，用户在同一次会话中的行为高度相关，在不同会话间的行为则相对独立。可以把用户行为日志按照时间间隔分割成会话并用SelfAttention机制建模它们之间的相互作用。我们试过又涨点了哈哈哈。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ca54a9",
   "metadata": {},
   "source": [
    "参考材料：\n",
    "\n",
    "* DIEN论文： https://arxiv.org/pdf/1809.03672.pdf \n",
    "\n",
    "* DIN+DIEN，机器学习唯一指定涨点技Attention： https://zhuanlan.zhihu.com/p/431131396\n",
    "\n",
    "* 从DIN到DIEN看阿里CTR算法的进化脉络： https://zhuanlan.zhihu.com/p/78365283\n",
    "\n",
    "* 代码实现参考： https://github.com/GitHub-HongweiZhang/prediction-flow\n",
    "\n",
    "上一篇文章我们介绍了DIN, 本篇文章我们介绍DIEN。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbead0b0",
   "metadata": {},
   "source": [
    "DIEN这篇文章的主要创新之处有3点：\n",
    "\n",
    "* 一是引入GRU来从用户行为日志序列中自然地抽取每个行为日志对应的用户兴趣表示(兴趣抽取层)。\n",
    "\n",
    "* 二是设计了一个辅助loss层，通过做一个辅助任务(区分真实的用户历史点击行为和负采样的非用户点击行为)来强化用户兴趣表示的学习。\n",
    "\n",
    "* 三是将注意力机制和GRU结构结合起来(AUGRU: Attention UPdate GRU)，来建模用户兴趣的时间演化得到最终的用户表示(兴趣演化层)。\n",
    "\n",
    "\n",
    "其中引入辅助Loss的技巧是神经网络涨点非常通用的一种高级技巧，值得我们学习。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dea5498",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4014082b",
   "metadata": {},
   "source": [
    "## 一，DIEN原理解析"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d06ec651",
   "metadata": {},
   "source": [
    "DIEN的主要出发点是，用户最近的行为可能比较远的行为更加重要。可以用循环神经网络GRU建模用户兴趣随时间的演化。\n",
    "\n",
    "DIEN选择的是不容易梯度消失且较快的GRU。\n",
    "\n",
    "![](https://tva1.sinaimg.cn/large/e6c9d24egy1h3x1brptqij20k10b8jsp.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3f7c443",
   "metadata": {},
   "source": [
    "### 1, 兴趣抽取层"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67b67c91",
   "metadata": {},
   "source": [
    "图中的 $b(t)$ 是用户的行为序列，而 $e(t)$是对应的embedding。随着自然发生的顺序， $e(t)$被输入GRU中，这就是兴趣抽取层。\n",
    "\n",
    "也是DIEN的第一条创新：引入GRU来从用户行为日志序列中自然地抽取每个行为日志对应的用户兴趣表示(兴趣抽取层)。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "459d6cd1",
   "metadata": {},
   "source": [
    "### 2，辅助loss "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a5afd9",
   "metadata": {},
   "source": [
    "如果忽略上面的AUGRU环节，GRU中的隐状态 $h(t)$就应该成为用户的行为序列最后的表示。\n",
    "\n",
    "如果直接就这样做，也不是不可以，但是$h(t)$学习到的东西可能不是我们想要的用户兴趣表示，或者说$h(t)$很难学习到有意义的信息。\n",
    "\n",
    "因为$h(t)$ 的迭代经过了很多步，然后还要和其他特征做拼接，然后还要经过MLP，最后才得到输出去计算Loss。\n",
    "\n",
    "这样的结果就是最后来了一个正样本或负样本，反向传播很难归因到 $h(t)$ 上。\n",
    "\n",
    "基于此DIEN给出了第二个要点：使用辅助Loss来强化$h(t)$的学习。\n",
    "\n",
    "我们来看看这个辅助Loss是怎么做的？这里设计了一个辅助任务，使用$h(t)$来区分真实的用户历史点击行为和负采样的非用户点击行为。\n",
    "\n",
    "由于$h(t)$ 代表着 t 时刻的用户兴趣表示，我们可以用它来预测 t+1时刻的广告用户是否点击。\n",
    "\n",
    "因为用户行为日志中都是用户点击过的广告(正样本, $e(t)$)，所以我们可以从全部的广告中给用户采样同样数量的用户没有点击过的广告作为负样本$e'(t)$。\n",
    "\n",
    "结合$h(t)$和 $e(t)$, $e'(t)$作为输入, 我们可以做一个二分类的辅助任务。\n",
    "\n",
    "这个辅助任务给$h(t)$在每个t时刻都提供了一个监督信号，使得$h(t)$能够更好地成为用户兴趣的抽取表示。\n",
    "\n",
    "真实应用场合下，你把开始的输入和最后的要求告诉网络，它就能给你一个好的结果的情况非常少。\n",
    "\n",
    "大多数时候是需要你去控制每一步的输入输出，每一步的loss才能防止网络各种偷懒作弊。\n",
    "\n",
    "辅助loss能够使得网络更受控制，向我们需要的方向发展，非常建议大家在实际业务中多试试辅助loss。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "525279a7",
   "metadata": {},
   "source": [
    "### 3，兴趣演化层"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d17c3da4",
   "metadata": {},
   "source": [
    "通过兴趣抽取层和辅助loss，我们得到了每个t时刻用户的一般兴趣表示。\n",
    "\n",
    "注意这个兴趣表示是一般性的，还没有和我们的候选广告做Attention关联。\n",
    "\n",
    "在DIN中，我们通过Attention机制构建了和候选广告相关的用户兴趣表示。\n",
    "\n",
    "而在DIEN中，我们希望建立的是和和候选广告相关，并且和时间演化相关的用户兴趣表示。\n",
    "\n",
    "DIEN通过结合Attention机制和GRU结构来做到这一点，这就是第三点创新AUGRU : Attention UPdate Gate GRU。\n",
    "\n",
    "下面我们进行详细讲解。\n",
    "\n",
    "\n",
    "一般地，各种RNN序列模型层(SimpleRNN,GRU,LSTM等)可以用函数表示如下:\n",
    "\n",
    "$$h_t = f(h_{t-1},i_t)$$\n",
    "\n",
    "这个公式的含义是：t时刻循环神经网络的输出向量$h_t$由t-1时刻的输出向量$h_{t-1}$和t时刻的输入$i_t$变换而来。\n",
    "\n",
    "为了结合Attention机制和GRU结构，我们需要设计这样的一个有三种输入的序列模型\n",
    "\n",
    "$$h_t = g(h_{t-1},i_t, a_t)$$\n",
    "\n",
    "这里的$a_t$是 t时刻的用户兴趣表示输入 $i_t$和候选广告计算出的attention 得分，是个标量。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf8b6b5e",
   "metadata": {},
   "source": [
    "我们先看看 GRU的 具体函数形式： \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "u_t &= \\sigma(W^u i_t + U^u h_{t-1} + b^u) \\tag{1} \\\\\n",
    "r_t &= \\sigma(W^r i_t + U^r h_{t-1} + b^r) \\tag{2} \\\\\n",
    "n_t &= \\tanh(W^n i_t + r_t \\circ U^n h_{t-1} + b^n) \\tag{3} \\\\\n",
    "h_t &= h_{t-1} - u_t \t\\circ h_{t-1} + u_t \\circ n_t \\tag{4} \\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38deb360",
   "metadata": {},
   "source": [
    "公式中的小圈表示哈达玛积，也就是两个向量逐位相乘。\n",
    "\n",
    "其中(1)式和(2)式计算的是更新门$u_t$和重置门$r_t$，是两个长度和$h_t$相同的向量。\n",
    "\n",
    "更新门用于控制每一步$h_t$被更新的比例，更新门越大，$h_t$更新幅度越大。\n",
    "\n",
    "重置门用于控制更新候选向量$n_t$中前一步的状态$h_{t-1}$被重新放入的比例，重置门越大，更新候选向量中$h_{t-1}$被重新放进来的比例越大。\n",
    "\n",
    "注意到(4)式 实际上和ResNet的残差结构是相似的，都是 f(x) = x + g(x) 的形式，可以有效地防止长序列学习反向传播过程中梯度消失问题。\n",
    "\n",
    "如何在GRU的基础上把attention得分融入进来呢？有以下一些非常自然的想法：\n",
    "\n",
    "* 1， 用$a_t$缩放输入$i_t$, 这就是AIGRU: Attention Input GRU。其含义是相关性高的在输入端进行放大。\n",
    "\n",
    "* 2， 用$a_t$代替GRU的更新门，这就是AGRU: Attention based GRU。其含义是用直接用相关性作为更新幅度。\n",
    "\n",
    "* 3， 用$a_t$缩放GRU的更新门$u_t$，这就是AUGRU:  Attention Update Gate GRU。其含义是用用相关性缩放更新幅度。\n",
    "\n",
    "\n",
    "AIGRU实际上并没有改变GRU的结构，只是改变了其输入，这种方式对Attention的使用比较含蓄，我把每个历史广告的相关性强弱通过输入告诉GRU，GRU你就给我好好学吧，希望你把相关性强的广告多长点到脑子里。但是这种方式效果不是很理想，即使是相关性为0的历史广告，也会对进行更新。\n",
    "\n",
    "AGRU是改变了GRU的结构的，并且对Attention的使用非常激进，完全删掉了GRU原有的的更新门，GRU你的脑子归Attention管了，遇到相关性高的广告，一定大大地记上一笔。不过AGRU也有一个缺陷，那就是Attention得分实际上是个标量，无法反应不同维度的差异。\n",
    "\n",
    "AUGRU也是改变了GRU的结构的，并且对Attention的使用比较折衷，让Attention缩放GRU原有的更新幅度。GRU我给你找了个搭档Attention，你更新前先问问它，你两一起决定该迈多大的步子吧。\n",
    "\n",
    "DIEN论文中通过对比实验发现AUGRU的效果最好。\n",
    "\n",
    "我们看看AUGRU的核心实现代码。基本上和公式是一致的，应用了F.linear函数来实现矩阵乘法和加偏置。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6916692",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from torch import nn \n",
    "\n",
    "class AttentionUpdateGateGRUCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, bias=True):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bias = bias\n",
    "        # (Wu|Wr|Wn)\n",
    "        self.weight_ih = nn.Parameter(\n",
    "            torch.Tensor(3 * hidden_size, input_size))\n",
    "        # (Uu|Ur|Un)\n",
    "        self.weight_hh = nn.Parameter(\n",
    "            torch.Tensor(3 * hidden_size, hidden_size))\n",
    "        if bias:\n",
    "            # (b_iu|b_ir|b_in)\n",
    "            self.bias_ih = nn.Parameter(torch.Tensor(3 * hidden_size))\n",
    "            # (b_hu|b_hr|b_hn)\n",
    "            self.bias_hh = nn.Parameter(torch.Tensor(3 * hidden_size))\n",
    "        else:\n",
    "            self.register_parameter('bias_ih', None)\n",
    "            self.register_parameter('bias_hh', None)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1.0 / (self.hidden_size)**0.5\n",
    "        for weight in self.parameters():\n",
    "            nn.init.uniform_(weight, -stdv, stdv)\n",
    "            \n",
    "    def forward(self, x, hx, att_score):\n",
    "        gi = F.linear(x, self.weight_ih, self.bias_ih)\n",
    "        gh = F.linear(hx, self.weight_hh, self.bias_hh)\n",
    "        i_r, i_u, i_n = gi.chunk(3, 1)\n",
    "        h_r, h_u, h_n = gh.chunk(3, 1)\n",
    "\n",
    "        resetgate = torch.sigmoid(i_r + h_r)\n",
    "        updategate = torch.sigmoid(i_u + h_u)\n",
    "        newgate = torch.tanh(i_n + resetgate * h_n)\n",
    "\n",
    "        updategate = att_score.view(-1, 1) * updategate\n",
    "        hy = (1-updategate)*hx +  updategate*newgate\n",
    "\n",
    "        return hy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b187e5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d761ca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ca041c0b",
   "metadata": {},
   "source": [
    "## 二，DIEN的pytorch实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b0a1ba9",
   "metadata": {},
   "source": [
    "下面是一个DIEN模型的完整pytorch实现。许多代码和DIN的实现是一样的。\n",
    "\n",
    "这里的AttentionGroup类用来建立候选广告属性，历史广告属性，以及负采样的广告属性的pair关系。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf9e25d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F \n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "from collections import OrderedDict\n",
    "\n",
    "class MaxPooling(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super(MaxPooling, self).__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, input):\n",
    "        return torch.max(input, self.dim)[0]\n",
    "\n",
    "\n",
    "class SumPooling(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super(SumPooling, self).__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, input):\n",
    "        return torch.sum(input, self.dim)\n",
    "\n",
    "class Dice(nn.Module):\n",
    "    \"\"\"\n",
    "    The Data Adaptive Activation Function in DIN, a generalization of PReLu.\n",
    "    \"\"\"\n",
    "    def __init__(self, emb_size, dim=2, epsilon=1e-8):\n",
    "        super(Dice, self).__init__()\n",
    "        assert dim == 2 or dim == 3\n",
    "\n",
    "        self.bn = nn.BatchNorm1d(emb_size, eps=epsilon)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        self.dim = dim\n",
    "        \n",
    "        # wrap alpha in nn.Parameter to make it trainable\n",
    "        self.alpha = nn.Parameter(torch.zeros((emb_size,))) if self.dim == 2 else nn.Parameter(\n",
    "            torch.zeros((emb_size, 1)))\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        assert x.dim() == self.dim\n",
    "        if self.dim == 2:\n",
    "            x_p = self.sigmoid(self.bn(x))\n",
    "            out = self.alpha * (1 - x_p) * x + x_p * x\n",
    "        else:\n",
    "            x = torch.transpose(x, 1, 2)\n",
    "            x_p = self.sigmoid(self.bn(x))\n",
    "            out = self.alpha * (1 - x_p) * x + x_p * x\n",
    "            out = torch.transpose(out, 1, 2)\n",
    "        return out\n",
    "\n",
    "    \n",
    "class Identity(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    def forward(self, x):\n",
    "        return x\n",
    "    \n",
    "def get_activation_layer(name, hidden_size=None, dice_dim=2):\n",
    "    name = name.lower()\n",
    "    name_dict = {x.lower():x for x in dir(nn) if '__' not in x and 'Z'>=x[0]>='A'}\n",
    "    if name==\"linear\":\n",
    "        return Identity()\n",
    "    elif name==\"dice\":\n",
    "        assert dice_dim\n",
    "        return Dice(hidden_size, dice_dim)\n",
    "    else:\n",
    "        assert name in name_dict, f'activation type {name} not supported!'\n",
    "        return getattr(nn,name_dict[name])()\n",
    "    \n",
    "def init_weights(model):\n",
    "    if isinstance(model, nn.Linear):\n",
    "        if model.weight is not None:\n",
    "            nn.init.kaiming_uniform_(model.weight.data)\n",
    "        if model.bias is not None:\n",
    "            nn.init.normal_(model.bias.data)\n",
    "    elif isinstance(model, (nn.BatchNorm1d,nn.BatchNorm2d,nn.BatchNorm3d)):\n",
    "        if model.weight is not None:\n",
    "            nn.init.normal_(model.weight.data, mean=1, std=0.02)\n",
    "        if model.bias is not None:\n",
    "            nn.init.constant_(model.bias.data, 0)\n",
    "    else:\n",
    "        pass\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, input_size, hidden_layers,\n",
    "                 dropout=0.0, batchnorm=True, activation='relu'):\n",
    "        super(MLP, self).__init__()\n",
    "        modules = OrderedDict()\n",
    "        previous_size = input_size\n",
    "        for index, hidden_layer in enumerate(hidden_layers):\n",
    "            modules[f\"dense{index}\"] = nn.Linear(previous_size, hidden_layer)\n",
    "            if batchnorm:\n",
    "                modules[f\"batchnorm{index}\"] = nn.BatchNorm1d(hidden_layer)\n",
    "            if activation:\n",
    "                modules[f\"activation{index}\"] = get_activation_layer(activation,hidden_layer,2)\n",
    "            if dropout:\n",
    "                modules[f\"dropout{index}\"] = nn.Dropout(dropout)\n",
    "            previous_size = hidden_layer\n",
    "        self.mlp = nn.Sequential(modules)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mlp(x)\n",
    "\n",
    "\n",
    "class AttentionGRUCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, bias=True):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bias = bias\n",
    "        # (Wr|Wn)\n",
    "        self.weight_ih = nn.Parameter(\n",
    "            torch.Tensor(2 * hidden_size, input_size))\n",
    "        # (Ur|Un)\n",
    "        self.weight_hh = nn.Parameter(\n",
    "            torch.Tensor(2 * hidden_size, hidden_size))\n",
    "        if bias:\n",
    "            # (b_ir|b_in)\n",
    "            self.bias_ih = nn.Parameter(torch.Tensor(2 * hidden_size))\n",
    "            # (b_hr|b_hn)\n",
    "            self.bias_hh = nn.Parameter(torch.Tensor(2 * hidden_size))\n",
    "        else:\n",
    "            self.register_parameter('bias_ih', None)\n",
    "            self.register_parameter('bias_hh', None)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1.0 / (self.hidden_size)**0.5\n",
    "        for weight in self.parameters():\n",
    "            nn.init.uniform_(weight, -stdv, stdv)\n",
    "\n",
    "    def forward(self, x, hx, att_score):\n",
    "\n",
    "        gi = F.linear(x, self.weight_ih, self.bias_ih)\n",
    "        gh = F.linear(hx, self.weight_hh, self.bias_hh)\n",
    "        i_r, i_n = gi.chunk(2, 1)\n",
    "        h_r, h_n = gh.chunk(2, 1)\n",
    "\n",
    "        resetgate = torch.sigmoid(i_r + h_r)\n",
    "        newgate = torch.tanh(i_n + resetgate * h_n)\n",
    "        att_score = att_score.view(-1, 1)\n",
    "        hy = (1. - att_score) * hx + att_score * newgate\n",
    "        \n",
    "        return hy\n",
    "\n",
    "\n",
    "class AttentionUpdateGateGRUCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, bias=True):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bias = bias\n",
    "        # (Wu|Wr|Wn)\n",
    "        self.weight_ih = nn.Parameter(\n",
    "            torch.Tensor(3 * hidden_size, input_size))\n",
    "        # (Uu|Ur|Un)\n",
    "        self.weight_hh = nn.Parameter(\n",
    "            torch.Tensor(3 * hidden_size, hidden_size))\n",
    "        if bias:\n",
    "            # (b_iu|b_ir|b_in)\n",
    "            self.bias_ih = nn.Parameter(torch.Tensor(3 * hidden_size))\n",
    "            # (b_hu|b_hr|b_hn)\n",
    "            self.bias_hh = nn.Parameter(torch.Tensor(3 * hidden_size))\n",
    "        else:\n",
    "            self.register_parameter('bias_ih', None)\n",
    "            self.register_parameter('bias_hh', None)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1.0 / (self.hidden_size)**0.5\n",
    "        for weight in self.parameters():\n",
    "            nn.init.uniform_(weight, -stdv, stdv)\n",
    "            \n",
    "    def forward(self, x, hx, att_score):\n",
    "        gi = F.linear(x, self.weight_ih, self.bias_ih)\n",
    "        gh = F.linear(hx, self.weight_hh, self.bias_hh)\n",
    "        i_u,i_r, i_n = gi.chunk(3, 1)\n",
    "        h_u,h_r, h_n = gh.chunk(3, 1)\n",
    "\n",
    "        updategate = torch.sigmoid(i_u + h_u)\n",
    "        resetgate = torch.sigmoid(i_r + h_r)\n",
    "        newgate = torch.tanh(i_n + resetgate * h_n)\n",
    "\n",
    "        updategate = att_score.view(-1, 1) * updategate\n",
    "        hy = (1-updategate)*hx +  updategate*newgate\n",
    "\n",
    "        return hy\n",
    "\n",
    "\n",
    "\n",
    "class DynamicGRU(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, bias=True, gru_type='AGRU'):\n",
    "        super(DynamicGRU, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        if gru_type == 'AGRU':\n",
    "            self.rnn = AttentionGRUCell(input_size, hidden_size, bias)\n",
    "        elif gru_type == 'AUGRU':\n",
    "            self.rnn = AttentionUpdateGateGRUCell(\n",
    "                input_size, hidden_size, bias)\n",
    "\n",
    "    def forward(self, x, att_scores, hx=None):\n",
    "        is_packed_input = isinstance(x, nn.utils.rnn.PackedSequence)\n",
    "        if not is_packed_input:\n",
    "            raise NotImplementedError(\n",
    "                \"DynamicGRU only supports packed input\")\n",
    "\n",
    "        is_packed_att_scores = isinstance(att_scores, nn.utils.rnn.PackedSequence)\n",
    "        if not is_packed_att_scores:\n",
    "            raise NotImplementedError(\n",
    "                \"DynamicGRU only supports packed att_scores\")\n",
    "\n",
    "        x, batch_sizes, sorted_indices, unsorted_indices = x\n",
    "        att_scores, _, _, _ = att_scores\n",
    "\n",
    "        max_batch_size = batch_sizes[0]\n",
    "        max_batch_size = int(max_batch_size)\n",
    "\n",
    "        if hx is None:\n",
    "            hx = torch.zeros(\n",
    "                max_batch_size, self.hidden_size,\n",
    "                dtype=x.dtype, device=x.device)\n",
    "\n",
    "        outputs = torch.zeros(\n",
    "            x.size(0), self.hidden_size,\n",
    "            dtype=x.dtype, device=x.device)\n",
    "\n",
    "        begin = 0\n",
    "        for batch in batch_sizes:\n",
    "            new_hx = self.rnn(\n",
    "                x[begin: begin + batch],\n",
    "                hx[0:batch],\n",
    "                att_scores[begin: begin + batch])\n",
    "            outputs[begin: begin + batch] = new_hx\n",
    "            hx = new_hx\n",
    "            begin += batch\n",
    "\n",
    "        return nn.utils.rnn.PackedSequence(\n",
    "            outputs, batch_sizes, sorted_indices, unsorted_indices)\n",
    "    \n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            input_size,\n",
    "            hidden_layers,\n",
    "            dropout=0.0,\n",
    "            batchnorm=True,\n",
    "            activation='prelu',\n",
    "            return_scores=False):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.return_scores = return_scores\n",
    "        \n",
    "        self.mlp = MLP(\n",
    "            input_size=input_size * 4,\n",
    "            hidden_layers=hidden_layers,\n",
    "            dropout=dropout,\n",
    "            batchnorm=batchnorm,\n",
    "            activation=activation)\n",
    "        self.fc = nn.Linear(hidden_layers[-1], 1)\n",
    "\n",
    "    def forward(self, query, keys, keys_length):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        query: 2D tensor, [Batch, Hidden]\n",
    "        keys: 3D tensor, [Batch, Time, Hidden]\n",
    "        keys_length: 1D tensor, [Batch]\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        outputs: 2D tensor, [Batch, Hidden]\n",
    "        \"\"\"\n",
    "        batch_size, max_length, dim = keys.size()\n",
    "\n",
    "        query = query.unsqueeze(1).expand(-1, max_length, -1)\n",
    "\n",
    "        din_all = torch.cat(\n",
    "            [query, keys, query - keys, query * keys], dim=-1)\n",
    "\n",
    "        din_all = din_all.view(batch_size * max_length, -1)\n",
    "\n",
    "        outputs = self.mlp(din_all)\n",
    "\n",
    "        outputs = self.fc(outputs).view(batch_size, max_length)  # [B, T]\n",
    "\n",
    "        # Scale\n",
    "        outputs = outputs / (dim ** 0.5)\n",
    "\n",
    "        # Mask\n",
    "        mask = (torch.arange(max_length, device=keys_length.device).repeat(\n",
    "            batch_size, 1) < keys_length.view(-1, 1))\n",
    "        outputs[~mask] = -np.inf\n",
    "\n",
    "        # Activation\n",
    "        outputs = F.softmax(outputs, dim=1)  #DIN uses sigmoid,DIEN uses softmax; [B, T]\n",
    "\n",
    "        if not self.return_scores:\n",
    "            # Weighted sum\n",
    "            outputs = torch.matmul(\n",
    "                outputs.unsqueeze(1), keys).squeeze()  # [B, H]\n",
    "        return outputs \n",
    "    \n",
    "class AuxiliaryNet(nn.Module):\n",
    "    def __init__(self, input_size, hidden_layers, activation='sigmoid'):\n",
    "        super().__init__()\n",
    "        modules = OrderedDict()\n",
    "        previous_size = input_size\n",
    "        for index, hidden_layer in enumerate(hidden_layers):\n",
    "            modules[f\"dense{index}\"] = nn.Linear(previous_size, hidden_layer)\n",
    "            if activation:\n",
    "                modules[f\"activation{index}\"] = get_activation_layer(activation)\n",
    "            previous_size = hidden_layer\n",
    "        modules[\"final_layer\"] = nn.Linear(previous_size, 1)\n",
    "        self.mlp = nn.Sequential(modules)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.sigmoid(self.mlp(x))\n",
    "\n",
    "\n",
    "class Interest(nn.Module):\n",
    "    SUPPORTED_GRU_TYPE = ['GRU', 'AIGRU', 'AGRU', 'AUGRU']\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            input_size,\n",
    "            gru_type='AUGRU',\n",
    "            gru_dropout=0.0,\n",
    "            att_hidden_layers=[80, 40],\n",
    "            att_dropout=0.0,\n",
    "            att_batchnorm=True,\n",
    "            att_activation='prelu',\n",
    "            use_negsampling=False):\n",
    "        super(Interest, self).__init__()\n",
    "        if gru_type not in Interest.SUPPORTED_GRU_TYPE:\n",
    "            raise NotImplementedError(f\"gru_type: {gru_type} is not supported\")\n",
    "\n",
    "        self.gru_type = gru_type\n",
    "        self.use_negsampling = use_negsampling\n",
    "\n",
    "        self.interest_extractor = nn.GRU(\n",
    "            input_size=input_size,\n",
    "            hidden_size=input_size,\n",
    "            batch_first=True,\n",
    "            bidirectional=False)\n",
    "\n",
    "        if self.use_negsampling:\n",
    "            self.auxiliary_net = AuxiliaryNet(\n",
    "                input_size * 2, hidden_layers=[100, 50])\n",
    "\n",
    "        if gru_type == 'GRU':\n",
    "            self.attention = Attention(\n",
    "                input_size=input_size,\n",
    "                hidden_layers=att_hidden_layers,\n",
    "                dropout=att_dropout,\n",
    "                batchnorm=att_batchnorm,\n",
    "                activation=att_activation)\n",
    "            \n",
    "            self.interest_evolution = nn.GRU(\n",
    "                input_size=input_size,\n",
    "                hidden_size=input_size,\n",
    "                batch_first=True,\n",
    "                bidirectional=False)\n",
    "                \n",
    "        elif gru_type == 'AIGRU':\n",
    "            self.attention = Attention(\n",
    "                input_size=input_size,\n",
    "                hidden_layers=att_hidden_layers,\n",
    "                dropout=att_dropout,\n",
    "                batchnorm=att_batchnorm,\n",
    "                activation=att_activation,\n",
    "                return_scores=True)\n",
    "\n",
    "            self.interest_evolution = nn.GRU(\n",
    "                input_size=input_size,\n",
    "                hidden_size=input_size,\n",
    "                batch_first=True,\n",
    "                bidirectional=False)\n",
    "            \n",
    "        elif gru_type == 'AGRU' or gru_type == 'AUGRU':\n",
    "            self.attention = Attention(\n",
    "                input_size=input_size,\n",
    "                hidden_layers=att_hidden_layers,\n",
    "                dropout=att_dropout,\n",
    "                batchnorm=att_batchnorm,\n",
    "                activation=att_activation,\n",
    "                return_scores=True)\n",
    "\n",
    "            self.interest_evolution = DynamicGRU(\n",
    "                input_size=input_size,\n",
    "                hidden_size=input_size,\n",
    "                gru_type=gru_type)\n",
    "\n",
    "    @staticmethod\n",
    "    def get_last_state(states, keys_length):\n",
    "        # states [B, T, H]\n",
    "        batch_size, max_seq_length, hidden_size = states.size()\n",
    "\n",
    "        mask = (torch.arange(max_seq_length, device=keys_length.device).repeat(\n",
    "            batch_size, 1) == (keys_length.view(-1, 1) - 1))\n",
    "\n",
    "        return states[mask]\n",
    "\n",
    "    def cal_auxiliary_loss(\n",
    "            self, states, click_seq, noclick_seq, keys_length):\n",
    "        # states [B, T, H]\n",
    "        # click_seq [B, T, H]\n",
    "        # noclick_seq [B, T, H]\n",
    "        # keys_length [B]\n",
    "        batch_size, max_seq_length, embedding_size = states.size()\n",
    "\n",
    "        mask = (torch.arange(max_seq_length, device=states.device).repeat(\n",
    "            batch_size, 1) < keys_length.view(-1, 1)).float()\n",
    "\n",
    "        click_input = torch.cat([states, click_seq], dim=-1)\n",
    "        noclick_input = torch.cat([states, noclick_seq], dim=-1)\n",
    "        embedding_size = embedding_size * 2\n",
    "\n",
    "        click_p = self.auxiliary_net(\n",
    "            click_input.view(\n",
    "                batch_size * max_seq_length, embedding_size)).view(\n",
    "                    batch_size, max_seq_length)[mask > 0].view(-1, 1)\n",
    "        click_target = torch.ones(\n",
    "            click_p.size(), dtype=torch.float, device=click_p.device)\n",
    "\n",
    "        noclick_p = self.auxiliary_net(\n",
    "            noclick_input.view(\n",
    "                batch_size * max_seq_length, embedding_size)).view(\n",
    "                    batch_size, max_seq_length)[mask > 0].view(-1, 1)\n",
    "        noclick_target = torch.zeros(\n",
    "            noclick_p.size(), dtype=torch.float, device=noclick_p.device)\n",
    "\n",
    "        loss = F.binary_cross_entropy(\n",
    "            torch.cat([click_p, noclick_p], dim=0),\n",
    "            torch.cat([click_target, noclick_target], dim=0))\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def forward(self, query, keys, keys_length, neg_keys=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        query: 2D tensor, [Batch, Hidden]\n",
    "        keys: 3D tensor, [Batch, Time, Hidden]\n",
    "        keys_length: 1D tensor, [Batch]\n",
    "        neg_keys: 3D tensor, [Batch, Time, Hidden]\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        outputs: 2D tensor, [Batch, Hidden]\n",
    "        \"\"\"\n",
    "        batch_size, max_length, dim = keys.size()\n",
    "\n",
    "        packed_keys = pack_padded_sequence(\n",
    "            keys,\n",
    "            lengths=keys_length.squeeze().cpu(),\n",
    "            batch_first=True,\n",
    "            enforce_sorted=False)\n",
    "\n",
    "        packed_interests, _ = self.interest_extractor(packed_keys)\n",
    "\n",
    "        aloss = None\n",
    "        if (self.gru_type != 'GRU') or self.use_negsampling:\n",
    "            interests, _ = pad_packed_sequence(\n",
    "                packed_interests,\n",
    "                batch_first=True,\n",
    "                padding_value=0.0,\n",
    "                total_length=max_length)\n",
    "\n",
    "            if self.use_negsampling:\n",
    "                aloss = self.cal_auxiliary_loss(\n",
    "                    interests[:, :-1, :],\n",
    "                    keys[:, 1:, :],\n",
    "                    neg_keys[:, 1:, :],\n",
    "                    keys_length - 1)\n",
    "\n",
    "        if self.gru_type == 'GRU':\n",
    "            packed_interests, _ = self.interest_evolution(packed_interests)\n",
    "\n",
    "            interests, _ = pad_packed_sequence(\n",
    "                packed_interests,\n",
    "                batch_first=True,\n",
    "                padding_value=0.0,\n",
    "                total_length=max_length)\n",
    "\n",
    "            outputs = self.attention(query, interests, keys_length)\n",
    "\n",
    "        elif self.gru_type == 'AIGRU':\n",
    "            # attention\n",
    "            scores = self.attention(query, interests, keys_length)\n",
    "            interests = interests * scores.unsqueeze(-1)\n",
    "\n",
    "            packed_interests = pack_padded_sequence(\n",
    "                interests,\n",
    "                lengths=keys_length.squeeze().cpu(),\n",
    "                batch_first=True,\n",
    "                enforce_sorted=False)\n",
    "            _, outputs = self.interest_evolution(packed_interests)\n",
    "            outputs = outputs.squeeze()\n",
    "\n",
    "        elif self.gru_type == 'AGRU' or self.gru_type == 'AUGRU':\n",
    "            # attention\n",
    "            scores = self.attention(query, interests, keys_length)\n",
    "\n",
    "            packed_interests = pack_padded_sequence(\n",
    "                interests,\n",
    "                lengths=keys_length.squeeze().cpu(),\n",
    "                batch_first=True,\n",
    "                enforce_sorted=False)\n",
    "\n",
    "            packed_scores = pack_padded_sequence(\n",
    "                scores,\n",
    "                lengths=keys_length.squeeze().cpu(),\n",
    "                batch_first=True,\n",
    "                enforce_sorted=False)\n",
    "\n",
    "            outputs, _ = pad_packed_sequence(\n",
    "                self.interest_evolution(\n",
    "                    packed_interests, packed_scores), batch_first=True)\n",
    "            # pick last state\n",
    "            outputs = Interest.get_last_state(\n",
    "                outputs, keys_length.squeeze())\n",
    "\n",
    "        return outputs, aloss\n",
    "    \n",
    "class AttentionGroup(object):\n",
    "    def __init__(self, name, pairs,\n",
    "                 hidden_layers, activation='dice', att_dropout=0.0,\n",
    "                 gru_type='AUGRU', gru_dropout=0.0):\n",
    "        self.name = name\n",
    "        self.pairs = pairs\n",
    "        self.hidden_layers = hidden_layers\n",
    "        self.activation = activation\n",
    "        self.att_dropout = att_dropout\n",
    "        self.gru_type = gru_type\n",
    "        self.gru_dropout = gru_dropout\n",
    "\n",
    "        self.related_feature_names = set()\n",
    "        self.neg_feature_names = set()\n",
    "        for pair in pairs:\n",
    "            self.related_feature_names.add(pair['ad'])\n",
    "            self.related_feature_names.add(pair['pos_hist'])\n",
    "            if 'neg_hist' in pair:\n",
    "                self.related_feature_names.add(pair['neg_hist'])\n",
    "                self.neg_feature_names.add(pair['neg_hist'])\n",
    "\n",
    "    def is_attention_feature(self, feature_name):\n",
    "        if feature_name in self.related_feature_names:\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    def is_neg_sampling_feature(self, feature_name):\n",
    "        if feature_name in self.neg_feature_names:\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    @property\n",
    "    def pairs_count(self):\n",
    "        return len(self.pairs)\n",
    "    \n",
    "class DIEN(nn.Module):\n",
    "    def __init__(self, num_features,cat_features,seq_features, \n",
    "                 cat_nums,embedding_size, attention_groups,\n",
    "                 mlp_hidden_layers, mlp_activation='prelu', mlp_dropout=0.0,\n",
    "                 use_negsampling = False,\n",
    "                 d_out = 1\n",
    "                 ):\n",
    "        super().__init__()\n",
    "        self.num_features = num_features\n",
    "        self.cat_features = cat_features\n",
    "        self.seq_features = seq_features\n",
    "        self.cat_nums = cat_nums \n",
    "        self.embedding_size = embedding_size\n",
    "        \n",
    "        self.attention_groups = attention_groups\n",
    "        \n",
    "        self.mlp_hidden_layers = mlp_hidden_layers\n",
    "        self.mlp_activation = mlp_activation\n",
    "        self.mlp_dropout = mlp_dropout\n",
    "        \n",
    "        self.d_out = d_out\n",
    "        self.use_negsampling = use_negsampling\n",
    "        \n",
    "        #embedding\n",
    "        self.embeddings = OrderedDict()\n",
    "        for feature in self.cat_features+self.seq_features:\n",
    "            self.embeddings[feature] = nn.Embedding(\n",
    "                self.cat_nums[feature], self.embedding_size, padding_idx=0)\n",
    "            self.add_module(f\"embedding:{feature}\",self.embeddings[feature])\n",
    "\n",
    "        self.sequence_poolings = OrderedDict()\n",
    "        self.attention_poolings = OrderedDict()\n",
    "        total_embedding_sizes = 0\n",
    "        for feature in self.cat_features:\n",
    "            total_embedding_sizes += self.embedding_size\n",
    "        for feature in self.seq_features:\n",
    "            if not self.is_neg_sampling_feature(feature):\n",
    "                total_embedding_sizes += self.embedding_size\n",
    "        \n",
    "        #sequence_pooling\n",
    "        for feature in self.seq_features:\n",
    "            if not self.is_attention_feature(feature):\n",
    "                self.sequence_poolings[feature] = MaxPooling(1)\n",
    "                self.add_module(f\"pooling:{feature}\",self.sequence_poolings[feature])\n",
    "\n",
    "        #attention_pooling\n",
    "        for attention_group in self.attention_groups:\n",
    "            self.attention_poolings[attention_group.name] = (\n",
    "                self.create_attention_fn(attention_group))\n",
    "            self.add_module(f\"attention_pooling:{attention_group.name}\",\n",
    "                self.attention_poolings[attention_group.name])\n",
    "\n",
    "        total_input_size = total_embedding_sizes+len(self.num_features)\n",
    "        \n",
    "        self.mlp = MLP(\n",
    "            total_input_size,\n",
    "            mlp_hidden_layers,\n",
    "            dropout=mlp_dropout, batchnorm=True, activation=mlp_activation)\n",
    "        \n",
    "        self.final_layer = nn.Linear(mlp_hidden_layers[-1], self.d_out)\n",
    "        self.apply(init_weights)\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        final_layer_inputs = list()\n",
    "\n",
    "        # linear\n",
    "        number_inputs = list()\n",
    "        for feature in self.num_features:\n",
    "            number_inputs.append(x[feature].view(-1, 1))\n",
    "\n",
    "        embeddings = OrderedDict()\n",
    "        for feature in self.cat_features:\n",
    "            embeddings[feature] = self.embeddings[feature](x[feature])\n",
    "\n",
    "        for feature in self.seq_features:\n",
    "            if not self.is_attention_feature(feature):\n",
    "                embeddings[feature] = self.sequence_poolings[feature](\n",
    "                    self.embeddings[feature](x[feature]))\n",
    "\n",
    "        auxiliary_losses = []\n",
    "        for attention_group in self.attention_groups:\n",
    "            query = torch.cat(\n",
    "                [embeddings[pair['ad']]\n",
    "                 for pair in attention_group.pairs],\n",
    "                dim=-1)\n",
    "            pos_hist = torch.cat(\n",
    "                [self.embeddings[pair['pos_hist']](\n",
    "                    x[pair['pos_hist']]) for pair in attention_group.pairs],\n",
    "                dim=-1)\n",
    "            \n",
    "            #hist_length = torch.sum(hist>0,axis=1)\n",
    "            keys_length = torch.min(torch.cat(\n",
    "                [torch.sum(x[pair['pos_hist']]>0,axis=1).view(-1, 1)\n",
    "                 for pair in attention_group.pairs],\n",
    "                dim=-1), dim=-1)[0]\n",
    "    \n",
    "            neg_hist = None\n",
    "            if self.use_negsampling:\n",
    "                neg_hist = torch.cat(\n",
    "                    [self.embeddings[pair['neg_hist']](\n",
    "                        x[pair['neg_hist']])\n",
    "                     for pair in attention_group.pairs],\n",
    "                    dim=-1)\n",
    "                \n",
    "            embeddings[attention_group.name], tmp_loss = (\n",
    "                self.attention_poolings[attention_group.name](\n",
    "                    query, pos_hist, keys_length, neg_hist))\n",
    "            if tmp_loss is not None:\n",
    "                auxiliary_losses.append(tmp_loss)\n",
    "\n",
    "        emb_concat = torch.cat(number_inputs + [\n",
    "            emb for emb in embeddings.values()], dim=-1)\n",
    "\n",
    "        final_layer_inputs = self.mlp(emb_concat)\n",
    "\n",
    "        output = self.final_layer(final_layer_inputs)\n",
    "        \n",
    "        auxiliary_avg_loss = None\n",
    "        if auxiliary_losses:\n",
    "            auxiliary_avg_loss = auxiliary_losses[0]\n",
    "            size = len(auxiliary_losses)\n",
    "            for i in range(1, size):\n",
    "                auxiliary_avg_loss += auxiliary_losses[i]\n",
    "            auxiliary_avg_loss /= size\n",
    "            \n",
    "        if  self.d_out==1:\n",
    "            output = output.squeeze() \n",
    "            \n",
    "        return output, auxiliary_avg_loss\n",
    "\n",
    "    def create_attention_fn(self, attention_group):\n",
    "        return Interest(\n",
    "            attention_group.pairs_count * self.embedding_size,\n",
    "            gru_type=attention_group.gru_type,\n",
    "            gru_dropout=attention_group.gru_dropout,\n",
    "            att_hidden_layers=attention_group.hidden_layers,\n",
    "            att_dropout=attention_group.att_dropout,\n",
    "            att_activation=attention_group.activation,\n",
    "            use_negsampling=self.use_negsampling)\n",
    "    \n",
    "    def is_attention_feature(self, feature):\n",
    "        for group in self.attention_groups:\n",
    "            if group.is_attention_feature(feature):\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    def is_neg_sampling_feature(self, feature):\n",
    "        for group in self.attention_groups:\n",
    "            if group.is_neg_sampling_feature(feature):\n",
    "                return True\n",
    "        return False\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a262a70a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "dc902338",
   "metadata": {},
   "source": [
    "## 三，Movielens数据集完整范例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7bbdfc",
   "metadata": {},
   "source": [
    "下面是一个基于Movielens评价数据集的DIEN完整范例，根据用户过去对一些电影的评价结果，来预测用户对候选电影是否会给好评。\n",
    "\n",
    "这个数据集不大，用CPU就能跑。😁"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a31e866",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8af36ca2",
   "metadata": {},
   "source": [
    "### 1，准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8e07df",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import pandas as pd \n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.preprocessing import QuantileTransformer\n",
    "from sklearn.pipeline import Pipeline, FeatureUnion \n",
    "from sklearn.impute import SimpleImputer \n",
    "from collections import Counter\n",
    "\n",
    "class CategoryEncoder(BaseEstimator, TransformerMixin):\n",
    "    \n",
    "    def __init__(self, min_cnt=5, word2idx=None, idx2word=None):\n",
    "        super().__init__() \n",
    "        self.min_cnt = min_cnt\n",
    "        self.word2idx = word2idx if word2idx else dict()\n",
    "        self.idx2word = idx2word if idx2word else dict()\n",
    "\n",
    "    def fit(self, x, y=None):\n",
    "        if not self.word2idx:\n",
    "            counter = Counter(np.asarray(x).ravel())\n",
    "\n",
    "            selected_terms = sorted(\n",
    "                list(filter(lambda x: counter[x] >= self.min_cnt, counter)))\n",
    "\n",
    "            self.word2idx = dict(\n",
    "                zip(selected_terms, range(1, len(selected_terms) + 1)))\n",
    "            self.word2idx['__PAD__'] = 0\n",
    "            if '__UNKNOWN__' not in self.word2idx:\n",
    "                self.word2idx['__UNKNOWN__'] = len(self.word2idx)\n",
    "\n",
    "        if not self.idx2word:\n",
    "            self.idx2word = {\n",
    "                index: word for word, index in self.word2idx.items()}\n",
    "\n",
    "        return self\n",
    "\n",
    "    def transform(self, x):\n",
    "        transformed_x = list()\n",
    "        for term in np.asarray(x).ravel():\n",
    "            try:\n",
    "                transformed_x.append(self.word2idx[term])\n",
    "            except KeyError:\n",
    "                transformed_x.append(self.word2idx['__UNKNOWN__'])\n",
    "\n",
    "        return np.asarray(transformed_x, dtype=np.int64)\n",
    "\n",
    "    def dimension(self):\n",
    "        return len(self.word2idx)\n",
    "\n",
    "class SequenceEncoder(BaseEstimator, TransformerMixin):\n",
    "    def __init__(self, sep=' ', min_cnt=5, max_len=None,\n",
    "                 word2idx=None, idx2word=None):\n",
    "        super().__init__() \n",
    "        self.sep = sep\n",
    "        self.min_cnt = min_cnt\n",
    "        self.max_len = max_len\n",
    "\n",
    "        self.word2idx = word2idx if word2idx else dict()\n",
    "        self.idx2word = idx2word if idx2word else dict()\n",
    "\n",
    "    def fit(self, x, y=None):\n",
    "        if not self.word2idx:\n",
    "            counter = Counter()\n",
    "\n",
    "            max_len = 0\n",
    "            for sequence in np.array(x).ravel():\n",
    "                words = sequence.split(self.sep)\n",
    "                counter.update(words)\n",
    "                max_len = max(max_len, len(words))\n",
    "\n",
    "            if self.max_len is None:\n",
    "                self.max_len = max_len\n",
    "\n",
    "            # drop rare words\n",
    "            words = sorted(\n",
    "                list(filter(lambda x: counter[x] >= self.min_cnt, counter)))\n",
    "\n",
    "            self.word2idx = dict(zip(words, range(1, len(words) + 1)))\n",
    "            self.word2idx['__PAD__'] = 0\n",
    "            if '__UNKNOWN__' not in self.word2idx:\n",
    "                self.word2idx['__UNKNOWN__'] = len(self.word2idx)\n",
    "\n",
    "        if not self.idx2word:\n",
    "            self.idx2word = {\n",
    "                index: word for word, index in self.word2idx.items()}\n",
    "\n",
    "        if not self.max_len:\n",
    "            max_len = 0\n",
    "            for sequence in np.array(x).ravel():\n",
    "                words = sequence.split(self.sep)\n",
    "                max_len = max(max_len, len(words))\n",
    "            self.max_len = max_len\n",
    "\n",
    "        return self\n",
    "\n",
    "    def transform(self, x):\n",
    "        transformed_x = list()\n",
    "\n",
    "        for sequence in np.asarray(x).ravel():\n",
    "            words = list()\n",
    "            for word in sequence.split(self.sep):\n",
    "                try:\n",
    "                    words.append(self.word2idx[word])\n",
    "                except KeyError:\n",
    "                    words.append(self.word2idx['__UNKNOWN__'])\n",
    "\n",
    "            transformed_x.append(\n",
    "                np.asarray(words[0:self.max_len], dtype=np.int64))\n",
    "\n",
    "        return np.asarray(transformed_x, dtype=object)\n",
    "    \n",
    "    def dimension(self):\n",
    "        return len(self.word2idx)\n",
    "\n",
    "    def max_length(self):\n",
    "        return self.max_len\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcfccdb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91b4cf7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import QuantileTransformer\n",
    "from sklearn.pipeline import Pipeline \n",
    "from sklearn.impute import SimpleImputer \n",
    "from tqdm import tqdm \n",
    "\n",
    "dftrain = pd.read_csv(\"../data/ml_1m/train.csv\")\n",
    "dfval = pd.read_csv(\"../data/ml_1m/test.csv\")\n",
    "\n",
    "for col in [\"movieId\",\"histHighRatedMovieIds\",\"negHistMovieIds\",\"genres\"]:\n",
    "    dftrain[col] = dftrain[col].astype(str)\n",
    "    dfval[col] = dfval[col].astype(str)\n",
    "\n",
    "num_features = ['age']\n",
    "cat_features = ['gender', 'movieId', 'occupation', 'zipCode']\n",
    "seq_features = ['genres', 'histHighRatedMovieIds', 'negHistMovieIds']\n",
    "\n",
    "num_pipe = Pipeline(steps = [('impute',SimpleImputer()),('quantile',QuantileTransformer())])\n",
    "\n",
    "encoders = {}\n",
    "\n",
    "print(\"preprocess number features...\")\n",
    "dftrain[num_features] = num_pipe.fit_transform(dftrain[num_features]).astype(np.float32)\n",
    "dfval[num_features] = num_pipe.transform(dfval[num_features]).astype(np.float32)\n",
    "\n",
    "print(\"preprocess category features...\")\n",
    "for col in tqdm(cat_features):\n",
    "    encoders[col] = CategoryEncoder(min_cnt=5)\n",
    "    dftrain[col]  = encoders[col].fit_transform(dftrain[col])\n",
    "    dfval[col] =  encoders[col].transform(dfval[col])\n",
    "    \n",
    "print(\"preprocess sequence features...\")\n",
    "for col in tqdm(seq_features):\n",
    "    encoders[col] = SequenceEncoder(sep=\"|\",min_cnt=5)\n",
    "    dftrain[col]  = encoders[col].fit_transform(dftrain[col])\n",
    "    dfval[col] =  encoders[col].transform(dfval[col])\n",
    "    \n",
    "from collections import OrderedDict\n",
    "from itertools import chain\n",
    "from torch.utils.data import Dataset,DataLoader \n",
    "\n",
    "class Df2Dataset(Dataset):\n",
    "    def __init__(self, dfdata, num_features, cat_features,\n",
    "                 seq_features, encoders, label_col=\"label\"):\n",
    "        self.dfdata = dfdata\n",
    "        self.num_features = num_features\n",
    "        self.cat_features = cat_features \n",
    "        self.seq_features = seq_features\n",
    "        self.encoders = encoders\n",
    "        self.label_col = label_col\n",
    "        self.size = len(self.dfdata)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "    @staticmethod\n",
    "    def pad_sequence(sequence,max_length):\n",
    "        #zero is special index for padding\n",
    "        padded_seq = np.zeros(max_length, np.int32)\n",
    "        padded_seq[0: sequence.shape[0]] = sequence\n",
    "        return padded_seq\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        record = OrderedDict()\n",
    "        for col in self.num_features:\n",
    "            record[col] = self.dfdata[col].iloc[idx].astype(np.float32)\n",
    "            \n",
    "        for col in self.cat_features:\n",
    "            record[col] = self.dfdata[col].iloc[idx].astype(np.int64)\n",
    "            \n",
    "        for col in self.seq_features:\n",
    "            seq = self.dfdata[col].iloc[idx]\n",
    "            max_length = self.encoders[col].max_length()\n",
    "            record[col] = Df2Dataset.pad_sequence(seq,max_length)\n",
    "\n",
    "        if self.label_col is not None:\n",
    "            record['label'] = self.dfdata[self.label_col].iloc[idx].astype(np.float32)\n",
    "        return record\n",
    "\n",
    "    def get_num_batches(self, batch_size):\n",
    "        return np.ceil(self.size / batch_size)\n",
    "    \n",
    "ds_train = Df2Dataset(dftrain, num_features, cat_features, seq_features, encoders)\n",
    "ds_val = Df2Dataset(dfval,num_features, cat_features, seq_features, encoders)\n",
    "dl_train = DataLoader(ds_train, batch_size=128,shuffle=True)\n",
    "dl_val = DataLoader(ds_val,batch_size=128,shuffle=False)\n",
    "\n",
    "cat_nums = {k:v.dimension() for k,v in encoders.items()} \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf667338",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in dl_train:\n",
    "    break "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa47b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cat_nums)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4356ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7bca6ecc",
   "metadata": {},
   "source": [
    "### 2，定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80845bf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def create_net():\n",
    "    augru_attention_groups_with_neg = [\n",
    "    AttentionGroup(\n",
    "        name='group1',\n",
    "        pairs=[{'ad': 'movieId', 'pos_hist': 'histHighRatedMovieIds', 'neg_hist': 'negHistMovieIds'}],\n",
    "        hidden_layers=[16, 8], att_dropout=0.1, gru_type='AUGRU')\n",
    "    ]\n",
    "\n",
    "    net = DIEN(num_features=num_features,\n",
    "           cat_features=cat_features,\n",
    "           seq_features=seq_features,\n",
    "           cat_nums = cat_nums,\n",
    "           embedding_size=16,\n",
    "           attention_groups=augru_attention_groups_with_neg,\n",
    "           mlp_hidden_layers=[32,16],\n",
    "           mlp_activation=\"prelu\",\n",
    "           mlp_dropout=0.25,\n",
    "           use_negsampling=True,\n",
    "           d_out=1\n",
    "           )\n",
    "    \n",
    "    return net \n",
    "\n",
    "net = create_net() \n",
    "\n",
    "out,aloss = net.forward(batch)\n",
    "\n",
    "from torchkeras.summary import summary \n",
    "summary(net,input_data=batch);\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90f45a56",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a5731884",
   "metadata": {},
   "source": [
    "### 3，训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e2152d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b0f4f6",
   "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, batch):\n",
    "        #loss\n",
    "        preds,aloss = self.net(batch)\n",
    "        loss = self.loss_fn(preds,batch[\"label\"])+aloss\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, batch[\"label\"]).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",
    "            if self.stage==\"train\":\n",
    "                loss, step_metrics = self.steprunner(batch)\n",
    "            else:\n",
    "                with torch.no_grad():\n",
    "                    loss, step_metrics = self.steprunner(batch)\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)[0]\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",
    "                self.net.load_state_dict(torch.load(ckpt_path))\n",
    "                break \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",
    "        self.net.eval()\n",
    "        result = torch.cat([self.forward(t) for t in dataloader])\n",
    "        return result.data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b045f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f1cb716",
   "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",
    "#lr_scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.01, max_lr=0.1)\n",
    "\n",
    "model = KerasModel(net,\n",
    "                   loss_fn = loss_fn,\n",
    "                   metrics_dict= metrics_dict,\n",
    "                   optimizer = optimizer,\n",
    "                   #lr_scheduler = lr_scheduler\n",
    "                  )    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2770b9dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfhistory = model.fit(train_data=dl_train,val_data=dl_val,epochs=100, patience=10,\n",
    "                      monitor = \"val_auc\",mode=\"max\",ckpt_path='checkpoint.pt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2faeb1cb",
   "metadata": {},
   "source": [
    "![](https://tva1.sinaimg.cn/large/e6c9d24egy1h3z8ccgn9ij20p507qmy7.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca5d7aae",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "11fadf05",
   "metadata": {},
   "source": [
    "### 4，评估模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65a4c2a",
   "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": "43bca201",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(dfhistory,\"loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b680911e",
   "metadata": {},
   "source": [
    "![](https://tva1.sinaimg.cn/large/e6c9d24egy1h3z8afwf3zj20f20aiglx.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d33de6d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a104d3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(dfhistory,\"auc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94d485ac",
   "metadata": {},
   "source": [
    "![](https://tva1.sinaimg.cn/large/e6c9d24egy1h3z8ddym3bj20f20ab74o.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73dd35ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(dl_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d841d78d",
   "metadata": {},
   "source": [
    "{'val_loss': 0.7020544648170471, 'val_auc': 0.6469045281410217}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ba6eca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "648c8f9f",
   "metadata": {},
   "source": [
    "### 5，使用模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ba78fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score \n",
    "\n",
    "labels = torch.tensor([x[\"label\"] for x in ds_val])\n",
    "preds = model.predict(dl_val)\n",
    "val_auc = roc_auc_score(labels.cpu().numpy(),preds.cpu().numpy())\n",
    "print(val_auc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81e2621",
   "metadata": {},
   "source": [
    "```\n",
    "0.6469045283797497\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba29e7b8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "736439a5",
   "metadata": {},
   "source": [
    "### 6，保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "795aacc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.net.state_dict(),\"best_dien.pt\")\n",
    "net_clone = create_net()\n",
    "net_clone.load_state_dict(torch.load(\"best_dien.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec2eb79f",
   "metadata": {},
   "outputs": [],
   "source": [
    "net_clone.eval()\n",
    "labels = torch.tensor([x[\"label\"] for x in ds_val])\n",
    "preds = torch.cat([net_clone(x)[0].data for x in dl_val]) \n",
    "val_auc = roc_auc_score(labels.cpu().numpy(),preds.cpu().numpy())\n",
    "print(val_auc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bac02c9",
   "metadata": {},
   "source": [
    "```\n",
    "0.6469045283797497\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e26ac34",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f65cfe4b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md",
   "text_representation": {
    "extension": ".md",
    "format_name": "markdown",
    "format_version": "1.3",
    "jupytext_version": "1.11.2"
   }
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}