DSSM-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1ce23cc",
   "metadata": {},
   "source": [
    "# 导入所需的包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "95a3c75f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import datetime\n",
    "import time\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "#import faiss\n",
    "#from dnn_model import DNN\n",
    "#from encoder_model import Encoder\n",
    "#from weight_initiallizer import Initializer\n",
    "#from itemcf import itemcf_sim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5984a217",
   "metadata": {},
   "source": [
    "# 数据预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a7044a30",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_path = os.path.abspath('.')\n",
    "user = pd.read_csv(os.path.join(root_path, 'ml-1m', 'users.dat'), sep='::', names = ['user', 'gender', 'age', 'occupation', 'zip_code'],engine='python')\n",
    "movie = pd.read_csv(os.path.join(root_path, 'ml-1m', 'movies.dat'), sep='::', names = ['movie', 'title', 'genres'],engine='python')\n",
    "rating = pd.read_csv(os.path.join(root_path, 'ml-1m', 'ratings.dat'), sep='::', names = ['user', 'movie', 'ratings', 'timestamp'],engine='python')\n",
    "# mapping\n",
    "user_id_dict = dict()\n",
    "for idx, uid in enumerate(user['user'].tolist()):\n",
    "    user_id_dict[uid] = idx\n",
    "movie_id_dict = dict()\n",
    "for idx, mid in enumerate(movie['movie'].tolist()):\n",
    "    movie_id_dict[mid] = idx\n",
    "user['user'] = user['user'].map(user_id_dict)\n",
    "movie['movie'] = movie['movie'].map(movie_id_dict)\n",
    "rating['user'] = rating['user'].map(user_id_dict)\n",
    "rating['movie'] = rating['movie'].map(movie_id_dict)\n",
    "\n",
    "\n",
    "# 时间处理\n",
    "rating['timestamp'] = rating['timestamp'].apply(lambda x: time.localtime(x))\n",
    "rating['time_str'] = rating['timestamp'].apply(lambda x: \n",
    "                                                  time.strftime(\"%Y-%m-%d %H:%M:%S\",x))\n",
    "\n",
    "# 切分训练和验证集，数据是2000到2003年的，以2003年用户的行为为验证集，用用户2000-2002年数据进行预测和验证，过滤掉那些只在2003年有行为的用户\n",
    "train_data = rating[rating['time_str']<'2003-01-01 00:00:00']\n",
    "val_data = rating[rating['time_str']>='2003-01-01 00:00:00']\n",
    "# 过滤\n",
    "val_data = val_data[val_data['user'].isin(train_data['user'].unique())]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8c241d7",
   "metadata": {},
   "source": [
    "## DSSM loss 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1e7f1519",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 996861/996861 [08:48<00:00, 1887.19it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user</th>\n",
       "      <th>movie</th>\n",
       "      <th>tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1176</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>2698</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>1910</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>3121</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>655</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user  movie  tag\n",
       "0     0   1176    1\n",
       "1     0   2698    0\n",
       "2     0   1910    0\n",
       "3     0   3121    0\n",
       "4     0    655    1"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 随机负采样\n",
    "sample_list = list(train_data['movie'].unique())\n",
    "data = list()\n",
    "for idx, rows in tqdm(train_data.iterrows(), total=len(train_data)):\n",
    "    use = rows['user']\n",
    "    mov = rows['movie']\n",
    "    data.append([use, mov, 1])\n",
    "    for m in np.random.choice(sample_list, 3):\n",
    "        data.append([use, m, 0])\n",
    "data = pd.DataFrame(data, columns=['user', 'movie', 'tag'])\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4865122f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user</th>\n",
       "      <th>movie</th>\n",
       "      <th>tag</th>\n",
       "      <th>gender</th>\n",
       "      <th>age</th>\n",
       "      <th>occupation</th>\n",
       "      <th>genres</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1176</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>[16]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>2698</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>[11]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>1910</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>[10]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>3121</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>[1, 15]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>655</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>[3, 13, 14]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user  movie  tag  gender  age  occupation       genres\n",
       "0     0   1176    1       0    0          10         [16]\n",
       "1     0   2698    0       0    0          10         [11]\n",
       "2     0   1910    0       0    0          10         [10]\n",
       "3     0   3121    0       0    0          10      [1, 15]\n",
       "4     0    655    1       0    0          10  [3, 13, 14]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 合并用户特征和电影特征\n",
    "le = LabelEncoder()\n",
    "user['gender'] = le.fit_transform(user['gender'])\n",
    "user['age'] = le.fit_transform(user['age'])\n",
    "user['occupation'] = le.fit_transform(user['occupation'])\n",
    "\n",
    "data = pd.merge(data, user[['user', 'gender', 'age', 'occupation']], how='left', on='user')\n",
    "\n",
    "genres = list()\n",
    "tmp = movie['genres'].apply(lambda x: x.split('|'))\n",
    "for l in tmp.tolist():\n",
    "    genres += l\n",
    "    \n",
    "genres_dict = dict()\n",
    "for idx, g in  enumerate(list(set(genres))):\n",
    "    genres_dict[g] = idx + 1\n",
    "    \n",
    "movie['genres'] = tmp.apply(lambda x: [genres_dict[i] for i in x])\n",
    "\n",
    "data = pd.merge(data, movie[['movie', 'genres']], how='left', on='movie')\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f283392",
   "metadata": {},
   "source": [
    "# 定义训练网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1b2972f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class trainset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.x = data[0]\n",
    "        self.y = data[1]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        x = self.x[index]\n",
    "        y = self.y[index]\n",
    "        data = (x, y)\n",
    "        return data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.x)\n",
    "\n",
    "# 训练集验证集随机分割\n",
    "train_df, test_df = train_test_split(data, test_size=0.2, random_state=2021)\n",
    "train_x = train_df[['user', 'gender', 'age', 'occupation', 'movie']].values\n",
    "train_y = train_df['tag'].values\n",
    "test_x = test_df[['user', 'gender', 'age', 'occupation', 'movie']].values\n",
    "test_y = test_df['tag'].values\n",
    "\n",
    "# 构造dataloader\n",
    "train_dataset = trainset((train_x, train_y))\n",
    "test_dataset = trainset((test_x, test_y))\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91cf62a2",
   "metadata": {},
   "source": [
    "# DSSM MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "eccd7d09",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DNNModel(nn.Module):\n",
    "    def __init__(self, inp_user, inp_movie, out, input_user_categorical_feature, input_movie_categorical_feature, \n",
    "                 hidden_layers, dropouts, batch_norm):\n",
    "        super(DNNModel, self).__init__()\n",
    "        self.user_embed = nn.Embedding(input_user_categorical_feature[0][0], input_user_categorical_feature[0][1])\n",
    "        self.gender_embed = nn.Embedding(input_user_categorical_feature[1][0], input_user_categorical_feature[1][1])\n",
    "        self.age_embed = nn.Embedding(input_user_categorical_feature[2][0], input_user_categorical_feature[2][1])\n",
    "        self.occupation_embed = nn.Embedding(input_user_categorical_feature[3][0], input_user_categorical_feature[3][1])\n",
    "        self.movie_embed = nn.Embedding(input_movie_categorical_feature[0][0], input_movie_categorical_feature[0][1])\n",
    "#         self.genres_embed = nn.Embedding(input_movie_categorical_feature[1][0], input_movie_categorical_feature[1][1])\n",
    "        \n",
    "        self.user_dnn = nn.Sequential(\n",
    "            nn.Linear(512, 128),\n",
    "#             nn.Dropout(0.5),\n",
    "            nn.Linear(128, 64)\n",
    "        )\n",
    "        \n",
    "        self.movie_dnn = nn.Sequential(\n",
    "            nn.Linear(128, 128),\n",
    "#             nn.Dropout(0.5),\n",
    "            nn.Linear(128, 64)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        u = self.user_embed(x[:, 0])\n",
    "        g = self.gender_embed(x[:, 1])\n",
    "        a = self.age_embed(x[:, 2])\n",
    "        oc = self.occupation_embed(x[:, 3])\n",
    "        m = self.movie_embed(x[:, 4])\n",
    "        \n",
    "        u = torch.cat([u, g, a, oc], -1)\n",
    "        u = self.user_dnn(u)\n",
    "        m = self.movie_dnn(m)\n",
    "        u = u/torch.sum(u*u, 1).view(-1,1)\n",
    "        m = m/torch.sum(m*m, 1).view(-1,1)\n",
    "        return u, m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b70a8349",
   "metadata": {},
   "source": [
    "# train_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c99a18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_loader, val_loader, epoch, loss_function, optimizer, path, early_stop):\n",
    "    \"\"\"\n",
    "    pytorch 模型训练通用代码\n",
    "    :param model: pytorch 模型\n",
    "    :param train_loader: dataloader, 训练数据\n",
    "    :param val_loader: dataloader, 验证数据\n",
    "    :param epoch: int, 训练迭代次数\n",
    "    :param loss_function: 优化损失函数\n",
    "    :param optimizer: pytorch优化器\n",
    "    :param path: save path\n",
    "    :param early_stop: int, 提前停止步数\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "    # 是否使用GPU\n",
    "  #  device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    device = torch.device(\"cpu\")\n",
    "    model = model.to(device)\n",
    "    \n",
    "    # 多少步内验证集的loss没有变小就提前停止\n",
    "    patience, eval_loss = 0, 0\n",
    "    \n",
    "    # 训练\n",
    "    for i in range(epoch):\n",
    "        total_loss, count = 0, 0\n",
    "        y_pred = list()\n",
    "        y_true = list()\n",
    "        for idx, (x, y) in tqdm(enumerate(train_loader), total=len(train_loader)):\n",
    "            x, y = x.to(device), y.to(device) \n",
    "            u, m = model(x)\n",
    "            predict = torch.sigmoid(torch.sum(u*m, 1))\n",
    "            y_pred.extend(predict.cpu().detach().numpy())\n",
    "            y_true.extend(y.cpu().detach().numpy())\n",
    "            loss = loss_function(predict, y.float())\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += float(loss)\n",
    "            count += 1\n",
    "            \n",
    "        train_auc = roc_auc_score(np.array(y_true), np.array(y_pred))\n",
    "        torch.save(model, path.format(i+1))\n",
    "        print(\"Epoch %d train loss is %.3f and train auc is %.3f\" % (i+1, total_loss / count, train_auc))\n",
    "    \n",
    "        # 验证\n",
    "        total_eval_loss = 0\n",
    "        model.eval()\n",
    "        count_eval = 0\n",
    "        val_y_pred = list()\n",
    "        val_true = list()\n",
    "        for idx, (x, y) in tqdm(enumerate(val_loader), total=len(val_loader)):\n",
    "            x, y = x.to(device), y.to(device)\n",
    "            u, m = model(x)\n",
    "            predict = torch.sigmoid(torch.sum(u*m, 1))\n",
    "            val_y_pred.extend(predict.cpu().detach().numpy())\n",
    "            val_true.extend(y.cpu().detach().numpy())\n",
    "            loss = loss_function(predict, y.float())\n",
    "            total_eval_loss += float(loss)\n",
    "            count_eval += 1\n",
    "        val_auc = roc_auc_score(np.array(y_true), np.array(y_pred))\n",
    "        print(\"Epoch %d val loss is %.3fand train auc is %.3f\" % (i+1, total_eval_loss / count_eval, val_auc))\n",
    "        \n",
    "        # 提前停止策略\n",
    "        if i == 0:\n",
    "            eval_loss = total_eval_loss / count_eval\n",
    "        else:\n",
    "            if total_eval_loss / count_eval < eval_loss:\n",
    "                eval_loss = total_eval_loss / count_eval\n",
    "            else:\n",
    "                if patience < early_stop:\n",
    "                    patience += 1\n",
    "                else:\n",
    "                    print(\"val loss is not decrease in %d epoch and break training\" % patience)\n",
    "                    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "521ad96c",
   "metadata": {},
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a990530a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24922/24922 [13:29<00:00, 30.79it/s] \n",
      "  0%|          | 7/6231 [00:00<01:55, 53.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 train loss is 0.698 and train auc is 0.681\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6231/6231 [01:33<00:00, 66.42it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 val loss is 0.691and train auc is 0.681\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24922/24922 [24:03<00:00, 17.26it/s] \n",
      "  0%|          | 9/6231 [00:00<01:16, 81.21it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 train loss is 0.692 and train auc is 0.750\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6231/6231 [00:17<00:00, 355.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 val loss is 0.692and train auc is 0.750\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24922/24922 [08:08<00:00, 51.00it/s]\n",
      "  0%|          | 9/6231 [00:00<01:13, 84.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3 train loss is 0.691 and train auc is 0.737\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6231/6231 [00:17<00:00, 362.55it/s]\n",
      "  0%|          | 0/24922 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3 val loss is 0.690and train auc is 0.737\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24922/24922 [21:53<00:00, 18.98it/s] \n",
      "  0%|          | 14/6231 [00:00<00:44, 139.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4 train loss is 0.690 and train auc is 0.735\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6231/6231 [00:14<00:00, 443.23it/s]\n",
      "  0%|          | 0/24922 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4 val loss is 0.688and train auc is 0.735\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24922/24922 [23:20<00:00, 17.79it/s] \n",
      "  1%|          | 39/6231 [00:00<00:16, 384.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5 train loss is 0.690 and train auc is 0.737\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6231/6231 [00:14<00:00, 417.54it/s]\n",
      "  0%|          | 0/24922 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5 val loss is 0.688and train auc is 0.737\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 76%|███████▋  | 19008/24922 [07:12<01:55, 51.01it/s]"
     ]
    }
   ],
   "source": [
    "# 模型初始化\n",
    "inp_user = 128\n",
    "inp_movie = 128\n",
    "out = 64\n",
    "input_user_categorical_feature = {0: (6040, 128), 1: (2, 128), 2: (7, 128), 3: (21, 128)}\n",
    "input_movie_categorical_feature =  {0: (3883, 128), 1:(18, 128)}\n",
    "hidden_layers = [128, 64]\n",
    "dropouts = [0.5, 0.5, 0.5]\n",
    "batch_norm = False\n",
    "\n",
    "model = DNNModel(inp_user, inp_movie, out, input_user_categorical_feature, input_movie_categorical_feature, \n",
    "                 hidden_layers, dropouts, batch_norm)\n",
    "\n",
    "# 模型训练\n",
    "epoch = 20\n",
    "loss_function = F.binary_cross_entropy_with_logits\n",
    "early_stop = 3\n",
    "learn_rate = 0.004\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate)\n",
    "path = 'model/model_{}.pth'\n",
    "\n",
    "train_model(model, train_loader, test_loader, epoch, loss_function, optimizer, path, early_stop)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c359e207",
   "metadata": {},
   "source": [
    "验证结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe5b1a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 结果验证\n",
    "model.eval()\n",
    "user['movie'] = 1\n",
    "test_x = user[['user', 'gender', 'age', 'occupation', 'movie']].values\n",
    "x = torch.from_numpy(test_x).cuda()\n",
    "user_embed, _ = model(x)\n",
    "\n",
    "movie['user'] = 1\n",
    "movie['gender'] = 1\n",
    "movie['age'] = 1\n",
    "movie['occupation'] = 1\n",
    "test_x = movie[['user', 'gender', 'age', 'occupation','movie']].values\n",
    "x = torch.from_numpy(test_x).cuda()\n",
    "_, movie_embed = model(x)\n",
    "\n",
    "movie_embed = movie_embed.cpu().detach().numpy()\n",
    "user_embed = user_embed.cpu().detach().numpy()\n",
    "\n",
    "\n",
    "# faiss索引构建\n",
    "d = 64\n",
    "nlist = 10\n",
    "index = faiss.IndexFlatL2(d)\n",
    "index.add(movie_embed)\n",
    "\n",
    "# 验证集数据字典化\n",
    "user_movie_dict_val = dict()\n",
    "for idx, rows in tqdm(val_data.iterrows(), total=len(val_data)):\n",
    "    u = rows['user']\n",
    "    m = rows['movie']\n",
    "    if u not in user_movie_dict_val:\n",
    "        user_movie_dict_val[u] = [m]\n",
    "    else:\n",
    "         user_movie_dict_val[u].append(m)\n",
    "            \n",
    "# 用户推荐结果索引           \n",
    "D, I = index.search(user_embed[list(val_data['user'].unique())], 50)\n",
    "\n",
    "# 召回率计算\n",
    "hits, total = 0, 0\n",
    "for uid, rec_list in zip(list(val_data['user'].unique()), I):\n",
    "    hits += len(set(rec_list)&set(user_movie_dict_val[uid]))\n",
    "    total += len(user_movie_dict_val[uid])\n",
    "print(\"recall is %.3f\" % (hits/total))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07e1047",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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
}