2.5 Multitask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 多任务学习：在不同任务之间学习共性以及差异性，能够提高建模的质量以及效率。\n",
    "# 需要根据业务场景下多任务之间的相关性特点来设计和选型。警惕负迁移或跷跷板效应。\n",
    "\n",
    "# 共享方式：\n",
    "# Hard Parameter Sharing：底层是共享的隐藏层，学习各个任务的共同模式，上层用一些特定的全连接层学习特定任务模式\n",
    "# Soft Parameter Sharing：底层不使用共享的shared bottom，而是有多个tower，给不同的tower分配不同的权重。\n",
    "\n",
    "# 共享层次：\n",
    "# 共享多少部分的层，共享的是底层、中层还是高层？\n",
    "\n",
    "# MMOE、ESSM、CGC、PLE、kuaishouEBR、AITM\n",
    "# DSIN、AutoInt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "import gc\n",
    "# 加载数据集2\n",
    "# 数据集：ml-100k\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "# 加载数据\n",
    "ratings = np.array([[int(x) for x in line.strip().split('\\t')[:3]] for line in open('./data/ml-100k/ua.base','r').read().strip().split('\\n')], dtype=np.int32)\n",
    "ratings[:,-1] = (ratings[:,-1] - 0)/(max(ratings[:,-1]) - 0)\n",
    "occupation_dict = {'administrator':0, 'artist':1, 'doctor':2, 'educator':3, 'engineer':4, 'entertainment':5, 'executive':6, 'healthcare':7, 'homemaker':8, 'lawyer':9, 'librarian':10, 'marketing':11, 'none':12, 'other':13, 'programmer':14, 'retired':15, 'salesman':16, 'scientist':17, 'student':18, 'technician':19, 'writer':20}\n",
    "gender_dict={'M':1,'F':0}\n",
    "user_info = {}\n",
    "for line in open('./data/ml-100k/u.user','r', encoding='utf-8').read().strip().split('\\n'):\n",
    "    phs = line.strip().split('|')\n",
    "    user_info[int(phs[0])] = [int(phs[1]), gender_dict[phs[2]], occupation_dict[phs[3]]]\n",
    "item_info = {}\n",
    "for line in open('./data/ml-100k/u.item','r', encoding='ISO-8859-1').read().strip().split('\\n'):\n",
    "    phs = line.strip().split('|')\n",
    "    item_info[int(phs[0])] = phs[5:]\n",
    "# print(len(user_info[list(user_info.keys())[0]]), len(item_info[list(item_info.keys())[0]]))\n",
    "# print(data.shape)\n",
    "num_users = len(user_info)\n",
    "num_items = len(item_info)\n",
    "num_features = 22"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-07T03:41:23.464559500Z",
     "start_time": "2023-09-07T03:41:19.888271400Z"
    }
   },
   "id": "c5a7915fd42957cc"
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-09-05 16:37:53] epoch=[1/10], train_mse_overall_loss: 0.032886, train_mse_task1_loss: 0.007735, train_mse_task2_loss: 0.029018, validate_mse_overall_loss: 0.000093, validate_mse_task1_loss: 0.000021, validate_mse_task2_loss: 0.000083\n",
      "[2023-09-05 16:38:07] epoch=[2/10], train_mse_overall_loss: 0.000051, train_mse_task1_loss: 0.000012, train_mse_task2_loss: 0.000045, validate_mse_overall_loss: 0.000028, validate_mse_task1_loss: 0.000006, validate_mse_task2_loss: 0.000024\n",
      "[2023-09-05 16:38:20] epoch=[3/10], train_mse_overall_loss: 0.000019, train_mse_task1_loss: 0.000004, train_mse_task2_loss: 0.000017, validate_mse_overall_loss: 0.000013, validate_mse_task1_loss: 0.000003, validate_mse_task2_loss: 0.000011\n",
      "[2023-09-05 16:38:33] epoch=[4/10], train_mse_overall_loss: 0.000009, train_mse_task1_loss: 0.000002, train_mse_task2_loss: 0.000008, validate_mse_overall_loss: 0.000007, validate_mse_task1_loss: 0.000002, validate_mse_task2_loss: 0.000006\n",
      "[2023-09-05 16:38:47] epoch=[5/10], train_mse_overall_loss: 0.000005, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000005, validate_mse_overall_loss: 0.000004, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000004\n",
      "[2023-09-05 16:39:00] epoch=[6/10], train_mse_overall_loss: 0.000003, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000003, validate_mse_overall_loss: 0.000003, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000002\n",
      "[2023-09-05 16:39:13] epoch=[7/10], train_mse_overall_loss: 0.000002, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000002, validate_mse_overall_loss: 0.000002, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000002\n",
      "[2023-09-05 16:39:26] epoch=[8/10], train_mse_overall_loss: 0.000002, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 16:39:40] epoch=[9/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 16:39:53] epoch=[10/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000001\n"
     ]
    }
   ],
   "source": [
    "# ESSM:\n",
    "# Entire Space Multi-task Modeling via Post-Click Behavior Decomposition for Conversion Rate Prediction\n",
    "# 一个输入到两个相对独立网络，网络输出的概率按照其任务级联关系相乘。\n",
    "# soft sharing\n",
    "\n",
    "import os, gc\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Module, Parameter, MSELoss\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from torch.nn import Module, Sequential, ReLU, Dropout, Sigmoid\n",
    "\n",
    "# category: [1,2] + [all]\n",
    "# number: [0] + []\n",
    "number_feature_data = MinMaxScaler().fit_transform(np.array([[user_info[u][0]]  for u, i, r in ratings], dtype=np.float32))\n",
    "category_feature_data = np.array([user_info[u][1:] + item_info[i] for u, i, r in ratings], dtype=np.int32)\n",
    "data = np.concatenate([number_feature_data, category_feature_data, ratings[:,-1:]], axis=-1)\n",
    "num_number_features = number_feature_data.shape[-1]\n",
    "num_category_features = category_feature_data.shape[-1]\n",
    "num_features = data.shape[-1] - 1\n",
    "category_feature_vals = {}\n",
    "for i in range(num_number_features, num_features):\n",
    "    category_feature_vals[i] = sorted(list(set(list(data[:, i]))))\n",
    "    for rid in range(data.shape[0]):\n",
    "        data[rid, i] = category_feature_vals[i].index(data[rid, i])\n",
    "# print(len(user_info[list(user_info.keys())[0]]), len(item_info[list(item_info.keys())[0]]))\n",
    "# print(data.shape)\n",
    "\n",
    "num_users = len(user_info)\n",
    "num_items = len(item_info)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else ('mps:0' if torch.backends.mps.is_available() else \"cpu\"))\n",
    "batch_size = 100\n",
    "num_epochs = 10\n",
    "embedding_dim = 8 # sparse feature embedding dim\n",
    "X_train, X_test, y_train, y_test = train_test_split(data[:,:-1], data[:,-1], test_size=0.4, random_state=0)\n",
    "train_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train>1.8).float(), torch.from_numpy(y_train>3.8).float()), batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_test).float(), torch.from_numpy(y_test>1.8).float(), torch.from_numpy(y_test>3.8).float()), batch_size=batch_size, shuffle=False, pin_memory=True)\n",
    "\n",
    "class ESSM(Module):\n",
    "    def __init__(self, dense_feature_cols:[(int,int)], sparse_feature_cols:[(int,int)], sparse_feature_embedding_dim, \n",
    "                 num_task:int, dnn_layer_dims:list[int], dnn_dropout=0.):\n",
    "        super(ESSM, self).__init__()\n",
    "        self.dense_feature_cols, self.sparse_feature_cols, self.sparse_feature_embedding_dim = dense_feature_cols, sparse_feature_cols, sparse_feature_embedding_dim\n",
    "        self.num_task, self.dnn_layer_dims = num_task, dnn_layer_dims\n",
    "        # sparse feature embedding dict\n",
    "        self.embed_layers = nn.ModuleDict({'embed_' + str(i): nn.Embedding(num_embeddings=valcount, embedding_dim=sparse_feature_embedding_dim) for i, valcount in sparse_feature_cols})\n",
    "        self.input_dim = len(dense_feature_cols) + len(sparse_feature_cols) * sparse_feature_embedding_dim\n",
    "        # dnn for each task\n",
    "        self.dnn_nets = []\n",
    "        for i in range(num_task):\n",
    "            net = nn.Sequential()\n",
    "            pre_layer_dim = self.input_dim\n",
    "            for layer_dim in dnn_layer_dims:\n",
    "                net.append(nn.Linear(pre_layer_dim, layer_dim))\n",
    "                net.append(nn.BatchNorm1d(layer_dim))\n",
    "                net.append(nn.Dropout(dnn_dropout))\n",
    "                pre_layer_dim = layer_dim\n",
    "            net.append(nn.Linear(dnn_layer_dims[-1], 1))\n",
    "            net.append(nn.Sigmoid())\n",
    "            self.dnn_nets.append(net)\n",
    "    def forward(self, x):\n",
    "        dense_input = x[:, :len(self.dense_feature_cols)]\n",
    "        sparse_embeds = torch.cat([self.embed_layers['embed_' + str(i)](x[:, i].long()) for i in range(len(self.dense_feature_cols), x.shape[1])], axis=1)\n",
    "        x = torch.cat([sparse_embeds, dense_input], axis=-1)\n",
    "        outputs = [self.dnn_nets[i](x).squeeze() for i in range(self.num_task)]\n",
    "        # ESSM原文计算概率，然后相乘，我这里任务是评分预测，改成相加\n",
    "        output_sum = outputs[0]\n",
    "        for i in range(1,self.num_task):\n",
    "            outputs[i] = output_sum + outputs[i] # 这里注意不能用 +=，会inplace报错\n",
    "            output_sum = outputs[i]\n",
    "        return outputs\n",
    "    def parameters(self, recurse: bool = True):\n",
    "        paras = []\n",
    "        for layer in self.embed_layers.values():\n",
    "            paras += [p for p in layer.parameters()]\n",
    "        for net in self.dnn_nets:\n",
    "            paras += [p for p in net.parameters()]\n",
    "        return paras\n",
    "model = ESSM(dense_feature_cols=[i for i in range(num_number_features)], sparse_feature_cols=[(i,len(category_feature_vals[i])) for i in range(num_number_features, num_features)], sparse_feature_embedding_dim =embedding_dim, num_task=2, dnn_layer_dims=[128, 32], dnn_dropout=0.).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = MSELoss(reduction='sum').to(device)\n",
    "alpha = 0.5\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    # train:\n",
    "    epoch_train_losses = []\n",
    "    model.train()\n",
    "    for i, inputs in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1, norm_type=2)\n",
    "        optimizer.step()\n",
    "        epoch_train_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    # validate:\n",
    "    model.eval()\n",
    "    epoch_test_losses = []\n",
    "    for i, inputs in enumerate(test_loader):\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        epoch_test_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    num_train, num_test = sum([x[0] for x in epoch_train_losses]), sum([x[0] for x in epoch_test_losses])\n",
    "    train_overall_loss = sum([x[-1] for x in epoch_train_losses])/num_train\n",
    "    train_task1_loss = sum([x[1] for x in epoch_train_losses])/num_train\n",
    "    train_task2_loss = sum([x[2] for x in epoch_train_losses])/num_train\n",
    "    test_overall_loss = sum([x[-1] for x in epoch_test_losses])/num_test\n",
    "    test_task1_loss = sum([x[1] for x in epoch_test_losses])/num_test\n",
    "    test_task2_loss = sum([x[2] for x in epoch_test_losses])/num_test\n",
    "    # print\n",
    "    print('['+datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")+']', 'epoch=[{}/{}], train_mse_overall_loss: {:.6f}, train_mse_task1_loss: {:.6f}, train_mse_task2_loss: {:.6f}, validate_mse_overall_loss: {:.6f}, validate_mse_task1_loss: {:.6f}, validate_mse_task2_loss: {:.6f}'.format(epoch+1, num_epochs,  train_overall_loss, train_task1_loss, train_task2_loss, test_overall_loss, test_task1_loss, test_task2_loss))\n",
    "    gc.collect()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-05T08:39:53.829553100Z",
     "start_time": "2023-09-05T08:37:38.003601800Z"
    }
   },
   "id": "4a49ef56e4d88080"
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-09-05 16:40:36] epoch=[1/10], train_mse_overall_loss: 0.012091, train_mse_task1_loss: 0.007442, train_mse_task2_loss: 0.008370, validate_mse_overall_loss: 0.000036, validate_mse_task1_loss: 0.000023, validate_mse_task2_loss: 0.000025\n",
      "[2023-09-05 16:40:45] epoch=[2/10], train_mse_overall_loss: 0.000020, train_mse_task1_loss: 0.000013, train_mse_task2_loss: 0.000014, validate_mse_overall_loss: 0.000011, validate_mse_task1_loss: 0.000007, validate_mse_task2_loss: 0.000007\n",
      "[2023-09-05 16:40:54] epoch=[3/10], train_mse_overall_loss: 0.000007, train_mse_task1_loss: 0.000005, train_mse_task2_loss: 0.000005, validate_mse_overall_loss: 0.000005, validate_mse_task1_loss: 0.000003, validate_mse_task2_loss: 0.000003\n",
      "[2023-09-05 16:41:04] epoch=[4/10], train_mse_overall_loss: 0.000004, train_mse_task1_loss: 0.000002, train_mse_task2_loss: 0.000003, validate_mse_overall_loss: 0.000003, validate_mse_task1_loss: 0.000002, validate_mse_task2_loss: 0.000002\n",
      "[2023-09-05 16:41:12] epoch=[5/10], train_mse_overall_loss: 0.000002, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000002, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 16:41:21] epoch=[6/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 16:41:30] epoch=[7/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 16:41:39] epoch=[8/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-05 16:41:48] epoch=[9/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-05 16:41:57] epoch=[10/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n"
     ]
    }
   ],
   "source": [
    "# MMOE:\n",
    "# Modeling Task\n",
    "# 一个输入，多个输出。中间共享多个expert，并行每个任务在其上的gate，进入各自任务的tower网络，预测。模块概率共享。\n",
    "# soft sharing\n",
    "\n",
    "import os, gc\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Module, Parameter, MSELoss\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from torch.nn import Module, Sequential, ReLU, Dropout, Sigmoid\n",
    "\n",
    "# category: [1,2] + [all]\n",
    "# number: [0] + []\n",
    "number_feature_data = MinMaxScaler().fit_transform(np.array([[user_info[u][0]]  for u, i, r in ratings], dtype=np.float32))\n",
    "category_feature_data = np.array([user_info[u][1:] + item_info[i] for u, i, r in ratings], dtype=np.int32)\n",
    "data = np.concatenate([number_feature_data, category_feature_data, ratings[:,-1:]], axis=-1)\n",
    "num_number_features = number_feature_data.shape[-1]\n",
    "num_category_features = category_feature_data.shape[-1]\n",
    "num_features = data.shape[-1] - 1\n",
    "category_feature_vals = {}\n",
    "for i in range(num_number_features, num_features):\n",
    "    category_feature_vals[i] = sorted(list(set(list(data[:, i]))))\n",
    "    for rid in range(data.shape[0]):\n",
    "        data[rid, i] = category_feature_vals[i].index(data[rid, i])\n",
    "# print(len(user_info[list(user_info.keys())[0]]), len(item_info[list(item_info.keys())[0]]))\n",
    "# print(data.shape)\n",
    "\n",
    "num_users = len(user_info)\n",
    "num_items = len(item_info)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else ('mps:0' if torch.backends.mps.is_available() else \"cpu\"))\n",
    "batch_size = 100\n",
    "num_epochs = 10\n",
    "embedding_dim = 8 # sparse feature embedding dim\n",
    "X_train, X_test, y_train, y_test = train_test_split(data[:,:-1], data[:,-1], test_size=0.4, random_state=0)\n",
    "train_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train>1.8).float(), torch.from_numpy(y_train>3.8).float()), batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_test).float(), torch.from_numpy(y_test>1.8).float(), torch.from_numpy(y_test>3.8).float()), batch_size=batch_size, shuffle=False, pin_memory=True)\n",
    "\n",
    "class MMOE(Module):\n",
    "    def __init__(self, dense_feature_cols:[(int,int)], sparse_feature_cols:[(int,int)], sparse_feature_embedding_dim, \n",
    "                 hidden_dim:int, num_task:int, n_expert:int, dnn_layer_dims:list[int], dnn_dropout=0.):\n",
    "        super(MMOE, self).__init__()\n",
    "        self.dense_feature_cols, self.sparse_feature_cols, self.sparse_feature_embedding_dim = dense_feature_cols, sparse_feature_cols, sparse_feature_embedding_dim\n",
    "        self.num_task, self.n_expert, self.dnn_layer_dims, self.hidden_dim = num_task, n_expert, dnn_layer_dims, hidden_dim\n",
    "        # sparse feature embedding dict\n",
    "        self.embed_layers = nn.ModuleDict({'embed_' + str(i): nn.Embedding(num_embeddings=valcount, embedding_dim=sparse_feature_embedding_dim) for i, valcount in sparse_feature_cols})\n",
    "        self.input_dim = len(dense_feature_cols) + len(sparse_feature_cols) * sparse_feature_embedding_dim\n",
    "        # experts: 线性层\n",
    "        self.experts = torch.nn.Parameter(torch.randn(self.input_dim, hidden_dim, n_expert), requires_grad=True)\n",
    "        self.experts_bias = torch.nn.Parameter(torch.randn(hidden_dim, n_expert), requires_grad=True)\n",
    "        # gates: 线性层\n",
    "        self.gates = [torch.nn.Parameter(torch.randn(self.input_dim, n_expert), requires_grad=True) for i in range(num_task)]\n",
    "        self.gates_bias = [torch.nn.Parameter(torch.randn(n_expert), requires_grad=True) for i in range(num_task)]\n",
    "        # dnn for each task\n",
    "        self.dnn_nets = []\n",
    "        for i in range(num_task):\n",
    "            net = nn.Sequential()\n",
    "            pre_layer_dim = hidden_dim\n",
    "            for layer_dim in dnn_layer_dims:\n",
    "                net.append(nn.Linear(pre_layer_dim, layer_dim))\n",
    "                net.append(nn.BatchNorm1d(layer_dim))\n",
    "                net.append(nn.Dropout(dnn_dropout))\n",
    "                pre_layer_dim = layer_dim\n",
    "            net.append(nn.Linear(dnn_layer_dims[-1], 1))\n",
    "            net.append(nn.Sigmoid())\n",
    "            self.dnn_nets.append(net)\n",
    "    def forward(self, x):\n",
    "        dense_input = x[:, :len(self.dense_feature_cols)]\n",
    "        sparse_embeds = torch.cat([self.embed_layers['embed_' + str(i)](x[:, i].long()) for i in range(len(self.dense_feature_cols), x.shape[1])], axis=1)\n",
    "        x = torch.cat([sparse_embeds, dense_input], axis=-1)\n",
    "        experts_out = torch.einsum('ij, jkl -> ikl', x, self.experts) + self.experts_bias.unsqueeze(0) # [batch, hidden_dim, n_expert]\n",
    "        gates_out = [torch.softmax(torch.einsum('ab, bc -> ac', x, gate) + self.gates_bias[i].unsqueeze(0), dim=-1) for i, gate in enumerate(self.gates)] # [batch, n_expert]\n",
    "        towers_input = [torch.bmm(experts_out, gate_out.unsqueeze(-1)).squeeze() for gate_out in gates_out]\n",
    "        outputs = [self.dnn_nets[i](tower_input).squeeze() for i,tower_input in enumerate(towers_input)]\n",
    "        return outputs\n",
    "    def parameters(self, recurse: bool = True):\n",
    "        paras = [self.experts, self.experts_bias] + self.gates + self.gates_bias\n",
    "        for layer in self.embed_layers.values():\n",
    "            paras += [p for p in layer.parameters()]\n",
    "        for net in self.dnn_nets:\n",
    "            paras += [p for p in net.parameters()]\n",
    "        return paras\n",
    "model = MMOE(dense_feature_cols=[i for i in range(num_number_features)], sparse_feature_cols=[(i,len(category_feature_vals[i])) for i in range(num_number_features, num_features)], sparse_feature_embedding_dim =embedding_dim, hidden_dim=embedding_dim, num_task=2, n_expert=3, dnn_layer_dims=[128, 32], dnn_dropout=0.).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = MSELoss(reduction='sum').to(device)\n",
    "alpha = 0.5\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    # train:\n",
    "    epoch_train_losses = []\n",
    "    model.train()\n",
    "    for i, inputs in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1, norm_type=2)\n",
    "        optimizer.step()\n",
    "        epoch_train_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    # validate:\n",
    "    model.eval()\n",
    "    epoch_test_losses = []\n",
    "    for i, inputs in enumerate(test_loader):\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        epoch_test_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    num_train, num_test = sum([x[0] for x in epoch_train_losses]), sum([x[0] for x in epoch_test_losses])\n",
    "    train_overall_loss = sum([x[-1] for x in epoch_train_losses])/num_train\n",
    "    train_task1_loss = sum([x[1] for x in epoch_train_losses])/num_train\n",
    "    train_task2_loss = sum([x[2] for x in epoch_train_losses])/num_train\n",
    "    test_overall_loss = sum([x[-1] for x in epoch_test_losses])/num_test\n",
    "    test_task1_loss = sum([x[1] for x in epoch_test_losses])/num_test\n",
    "    test_task2_loss = sum([x[2] for x in epoch_test_losses])/num_test\n",
    "    # print\n",
    "    print('['+datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")+']', 'epoch=[{}/{}], train_mse_overall_loss: {:.6f}, train_mse_task1_loss: {:.6f}, train_mse_task2_loss: {:.6f}, validate_mse_overall_loss: {:.6f}, validate_mse_task1_loss: {:.6f}, validate_mse_task2_loss: {:.6f}'.format(epoch+1, num_epochs,  train_overall_loss, train_task1_loss, train_task2_loss, test_overall_loss, test_task1_loss, test_task2_loss))\n",
    "    gc.collect()\n",
    "\n",
    "# 这种方法共享很稀疏，收敛很快"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-05T08:41:57.474389Z",
     "start_time": "2023-09-05T08:40:24.333973900Z"
    }
   },
   "id": "56b432fec9a7451c"
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-09-05 17:28:46] epoch=[1/10], train_mse_overall_loss: 0.011408, train_mse_task1_loss: 0.007582, train_mse_task2_loss: 0.007617, validate_mse_overall_loss: 0.000035, validate_mse_task1_loss: 0.000027, validate_mse_task2_loss: 0.000022\n",
      "[2023-09-05 17:28:57] epoch=[2/10], train_mse_overall_loss: 0.000019, train_mse_task1_loss: 0.000014, train_mse_task2_loss: 0.000012, validate_mse_overall_loss: 0.000010, validate_mse_task1_loss: 0.000008, validate_mse_task2_loss: 0.000006\n",
      "[2023-09-05 17:29:11] epoch=[3/10], train_mse_overall_loss: 0.000007, train_mse_task1_loss: 0.000005, train_mse_task2_loss: 0.000004, validate_mse_overall_loss: 0.000005, validate_mse_task1_loss: 0.000004, validate_mse_task2_loss: 0.000003\n",
      "[2023-09-05 17:29:25] epoch=[4/10], train_mse_overall_loss: 0.000004, train_mse_task1_loss: 0.000003, train_mse_task2_loss: 0.000002, validate_mse_overall_loss: 0.000003, validate_mse_task1_loss: 0.000002, validate_mse_task2_loss: 0.000002\n",
      "[2023-09-05 17:29:36] epoch=[5/10], train_mse_overall_loss: 0.000002, train_mse_task1_loss: 0.000002, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000002, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 17:29:47] epoch=[6/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000001\n",
      "[2023-09-05 17:29:58] epoch=[7/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000001, train_mse_task2_loss: 0.000001, validate_mse_overall_loss: 0.000001, validate_mse_task1_loss: 0.000001, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-05 17:30:10] epoch=[8/10], train_mse_overall_loss: 0.000001, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-05 17:30:20] epoch=[9/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-05 17:30:31] epoch=[10/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n"
     ]
    }
   ],
   "source": [
    "# CGC:\n",
    "# 一个输入，多个输出。各自任务有自己的expert外，中间共享多个expert，并行每个任务在其上的gate，进入各自任务的tower网络，预测。模块概率共享。\n",
    "# 相较于MMOE，增加了任务私有的expert\n",
    "# soft sharing\n",
    "\n",
    "# PLE：是CGC叠加多个，然后进入各自任务的tower网络。（复现很简单，这里没有去实现）\n",
    "# Progressive layered extraction (ple): A novel multi-task learning (mtl) model for personalized recommendations， 2020\n",
    "\n",
    "import os, gc\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Module, Parameter, MSELoss\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from torch.nn import Module, Sequential, ReLU, Dropout, Sigmoid\n",
    "\n",
    "# category: [1,2] + [all]\n",
    "# number: [0] + []\n",
    "number_feature_data = MinMaxScaler().fit_transform(np.array([[user_info[u][0]]  for u, i, r in ratings], dtype=np.float32))\n",
    "category_feature_data = np.array([user_info[u][1:] + item_info[i] for u, i, r in ratings], dtype=np.int32)\n",
    "data = np.concatenate([number_feature_data, category_feature_data, ratings[:,-1:]], axis=-1)\n",
    "num_number_features = number_feature_data.shape[-1]\n",
    "num_category_features = category_feature_data.shape[-1]\n",
    "num_features = data.shape[-1] - 1\n",
    "category_feature_vals = {}\n",
    "for i in range(num_number_features, num_features):\n",
    "    category_feature_vals[i] = sorted(list(set(list(data[:, i]))))\n",
    "    for rid in range(data.shape[0]):\n",
    "        data[rid, i] = category_feature_vals[i].index(data[rid, i])\n",
    "# print(len(user_info[list(user_info.keys())[0]]), len(item_info[list(item_info.keys())[0]]))\n",
    "# print(data.shape)\n",
    "\n",
    "num_users = len(user_info)\n",
    "num_items = len(item_info)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else ('mps:0' if torch.backends.mps.is_available() else \"cpu\"))\n",
    "batch_size = 100\n",
    "num_epochs = 10\n",
    "embedding_dim = 8 # sparse feature embedding dim\n",
    "X_train, X_test, y_train, y_test = train_test_split(data[:,:-1], data[:,-1], test_size=0.4, random_state=0)\n",
    "train_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train>1.8).float(), torch.from_numpy(y_train>3.8).float()), batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_test).float(), torch.from_numpy(y_test>1.8).float(), torch.from_numpy(y_test>3.8).float()), batch_size=batch_size, shuffle=False, pin_memory=True)\n",
    "\n",
    "class Expert(Module):\n",
    "    def __init__(self, input_dim, hidden_dim:int, n_expert:int):\n",
    "        super(Expert, self).__init__()\n",
    "        self.input_dim, self.hidden_dim, self.n_expert = input_dim, hidden_dim, n_expert\n",
    "        # experts: 线性层\n",
    "        self.experts = torch.nn.Parameter(torch.randn(self.input_dim, hidden_dim, n_expert), requires_grad=True)\n",
    "        self.experts_bias = torch.nn.Parameter(torch.randn(hidden_dim, n_expert), requires_grad=True)\n",
    "    def forward(self, x):\n",
    "        expert_out = torch.einsum('ij, jkl -> ikl', x, self.experts) + self.experts_bias.unsqueeze(0) # [batch, hidden_dim, n_expert]\n",
    "        return expert_out.permute((0,2,1)) # [batch, n_expert, hidden_dim]\n",
    "    def parameters(self, recurse: bool = True):\n",
    "        return [self.experts, self.experts_bias]\n",
    "class CGC(Module):\n",
    "    def __init__(self, dense_feature_cols:[(int,int)], sparse_feature_cols:[(int,int)], sparse_feature_embedding_dim, \n",
    "                 hidden_dim:int, num_task:int, shared_n_expert:int, self_n_expert:int, dnn_layer_dims:list[int], dnn_dropout=0.):\n",
    "        super(CGC, self).__init__()\n",
    "        self.dense_feature_cols, self.sparse_feature_cols, self.sparse_feature_embedding_dim = dense_feature_cols, sparse_feature_cols, sparse_feature_embedding_dim\n",
    "        self.num_task, self.shared_n_expert, self.self_n_expert, self.dnn_layer_dims, self.hidden_dim = num_task, shared_n_expert, self_n_expert, dnn_layer_dims, hidden_dim\n",
    "        # sparse feature embedding dict\n",
    "        self.embed_layers = nn.ModuleDict({'embed_' + str(i): nn.Embedding(num_embeddings=valcount, embedding_dim=sparse_feature_embedding_dim) for i, valcount in sparse_feature_cols})\n",
    "        self.input_dim = len(dense_feature_cols) + len(sparse_feature_cols) * sparse_feature_embedding_dim\n",
    "        # experts\n",
    "        self.shared_expert = Expert(self.input_dim, hidden_dim, shared_n_expert)\n",
    "        self.self_experts = [Expert(self.input_dim, hidden_dim, self_n_expert) for i in range(num_task)]\n",
    "        # gates:\n",
    "        self.gates = [nn.Sequential(nn.Linear((shared_n_expert + self_n_expert) * hidden_dim, shared_n_expert + self_n_expert), nn.Softmax()) for i in range(num_task)]\n",
    "        # dnn for each task\n",
    "        self.dnn_nets = []\n",
    "        for i in range(num_task):\n",
    "            net = nn.Sequential()\n",
    "            pre_layer_dim = hidden_dim\n",
    "            for layer_dim in dnn_layer_dims:\n",
    "                net.append(nn.Linear(pre_layer_dim, layer_dim))\n",
    "                net.append(nn.BatchNorm1d(layer_dim))\n",
    "                net.append(nn.Dropout(dnn_dropout))\n",
    "                pre_layer_dim = layer_dim\n",
    "            net.append(nn.Linear(dnn_layer_dims[-1], 1))\n",
    "            net.append(nn.Sigmoid())\n",
    "            self.dnn_nets.append(net)\n",
    "    def forward(self, x):\n",
    "        batch_len = x.shape[0]\n",
    "        dense_input = x[:, :len(self.dense_feature_cols)]\n",
    "        sparse_embeds = torch.cat([self.embed_layers['embed_' + str(i)](x[:, i].long()) for i in range(len(self.dense_feature_cols), x.shape[1])], axis=1)\n",
    "        x = torch.cat([sparse_embeds, dense_input], axis=-1)\n",
    "        # torch.Size([100, 169]) torch.Size([100, 8, 3]) [torch.Size([100, 8, 2]), torch.Size([100, 8, 2])]\n",
    "        shared_expert_out = self.shared_expert(x)\n",
    "        experts_out = [torch.cat([shared_expert_out, self.self_experts[i](x)], dim=1) for i in range(self.num_task)]\n",
    "        gates_out = [gate(experts_out[i].reshape((batch_len,-1))) for i, gate in enumerate(self.gates)]\n",
    "        towers_input = [torch.bmm(experts_out[i].permute((0,2,1)), gate_out.unsqueeze(-1)).squeeze() for i, gate_out in enumerate(gates_out)]\n",
    "        outputs = [self.dnn_nets[i](tower_input).squeeze() for i,tower_input in enumerate(towers_input)]\n",
    "        return outputs\n",
    "    def parameters(self, recurse: bool = True):\n",
    "        paras = [p for p in self.shared_expert.parameters()]\n",
    "        for expert in self.self_experts:\n",
    "            paras += [p for p in expert.parameters()]\n",
    "        for gate in self.gates:\n",
    "            paras += [p for p in gate.parameters()]\n",
    "        for layer in self.embed_layers.values():\n",
    "            paras += [p for p in layer.parameters()]\n",
    "        for net in self.dnn_nets:\n",
    "            paras += [p for p in net.parameters()]\n",
    "        return paras\n",
    "model = CGC(dense_feature_cols=[i for i in range(num_number_features)], sparse_feature_cols=[(i,len(category_feature_vals[i])) for i in range(num_number_features, num_features)], sparse_feature_embedding_dim =embedding_dim, hidden_dim=embedding_dim, num_task=2, shared_n_expert=3, self_n_expert=2, dnn_layer_dims=[128, 32], dnn_dropout=0.).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = MSELoss(reduction='sum').to(device)\n",
    "alpha = 0.5\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    # train:\n",
    "    epoch_train_losses = []\n",
    "    model.train()\n",
    "    for i, inputs in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1, norm_type=2)\n",
    "        optimizer.step()\n",
    "        epoch_train_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    # validate:\n",
    "    model.eval()\n",
    "    epoch_test_losses = []\n",
    "    for i, inputs in enumerate(test_loader):\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        epoch_test_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    num_train, num_test = sum([x[0] for x in epoch_train_losses]), sum([x[0] for x in epoch_test_losses])\n",
    "    train_overall_loss = sum([x[-1] for x in epoch_train_losses])/num_train\n",
    "    train_task1_loss = sum([x[1] for x in epoch_train_losses])/num_train\n",
    "    train_task2_loss = sum([x[2] for x in epoch_train_losses])/num_train\n",
    "    test_overall_loss = sum([x[-1] for x in epoch_test_losses])/num_test\n",
    "    test_task1_loss = sum([x[1] for x in epoch_test_losses])/num_test\n",
    "    test_task2_loss = sum([x[2] for x in epoch_test_losses])/num_test\n",
    "    # print\n",
    "    print('['+datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")+']', 'epoch=[{}/{}], train_mse_overall_loss: {:.6f}, train_mse_task1_loss: {:.6f}, train_mse_task2_loss: {:.6f}, validate_mse_overall_loss: {:.6f}, validate_mse_task1_loss: {:.6f}, validate_mse_task2_loss: {:.6f}'.format(epoch+1, num_epochs,  train_overall_loss, train_task1_loss, train_task2_loss, test_overall_loss, test_task1_loss, test_task2_loss))\n",
    "    gc.collect()\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-05T09:30:31.199366800Z",
     "start_time": "2023-09-05T09:28:33.121316800Z"
    }
   },
   "id": "9d5eb170c07aefe0"
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-09-07 12:25:18] epoch=[1/10], train_mse_overall_loss: 0.004533, train_mse_task1_loss: 0.002647, train_mse_task2_loss: 0.003210, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:25:36] epoch=[2/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:25:53] epoch=[3/10], train_mse_overall_loss: 0.000004, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000004, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:26:11] epoch=[4/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:26:29] epoch=[5/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:26:46] epoch=[6/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:27:04] epoch=[7/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:27:21] epoch=[8/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:27:38] epoch=[9/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n",
      "[2023-09-07 12:27:56] epoch=[10/10], train_mse_overall_loss: 0.000000, train_mse_task1_loss: 0.000000, train_mse_task2_loss: 0.000000, validate_mse_overall_loss: 0.000000, validate_mse_task1_loss: 0.000000, validate_mse_task2_loss: 0.000000\n"
     ]
    }
   ],
   "source": [
    "# AITM: 建模漏斗任务中多层级级联关系，如点击 -> 购买\n",
    "# Modeling the Sequential Dependence among Audience Multi-step Conversions with Multi-task Learning in Targeted Display Advertising, KDD-2021\n",
    "# input embeddings -> 各自任务的tower 网络，低层任务的tower网络最后一层通过注意力机制(相当于自适应迁移信息)将信息融合到高层任务中（如上述点击任务到购买任务）\n",
    "# 高层桥接的注意力模式\n",
    "\n",
    "import os, gc\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Module, Parameter, MSELoss\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from torch.nn import Module, Sequential, ReLU, Dropout, Sigmoid\n",
    "\n",
    "# category: [1,2] + [all]\n",
    "# number: [0] + []\n",
    "number_feature_data = MinMaxScaler().fit_transform(np.array([[user_info[u][0]]  for u, i, r in ratings], dtype=np.float32))\n",
    "category_feature_data = np.array([user_info[u][1:] + item_info[i] for u, i, r in ratings], dtype=np.int32)\n",
    "data = np.concatenate([number_feature_data, category_feature_data, ratings[:,-1:]], axis=-1)\n",
    "num_number_features = number_feature_data.shape[-1]\n",
    "num_category_features = category_feature_data.shape[-1]\n",
    "num_features = data.shape[-1] - 1\n",
    "category_feature_vals = {}\n",
    "for i in range(num_number_features, num_features):\n",
    "    category_feature_vals[i] = sorted(list(set(list(data[:, i]))))\n",
    "    for rid in range(data.shape[0]):\n",
    "        data[rid, i] = category_feature_vals[i].index(data[rid, i])\n",
    "# print(len(user_info[list(user_info.keys())[0]]), len(item_info[list(item_info.keys())[0]]))\n",
    "# print(data.shape)\n",
    "\n",
    "num_users = len(user_info)\n",
    "num_items = len(item_info)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else ('mps:0' if torch.backends.mps.is_available() else \"cpu\"))\n",
    "batch_size = 100\n",
    "num_epochs = 10\n",
    "embedding_dim = 8 # sparse feature embedding dim\n",
    "X_train, X_test, y_train, y_test = train_test_split(data[:,:-1], data[:,-1], test_size=0.4, random_state=0)\n",
    "train_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train>1.8).float(), torch.from_numpy(y_train>3.8).float()), batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(dataset=TensorDataset(torch.from_numpy(X_test).float(), torch.from_numpy(y_test>1.8).float(), torch.from_numpy(y_test>3.8).float()), batch_size=batch_size, shuffle=False, pin_memory=True)\n",
    "\n",
    "# 双向迁移模块\n",
    "class AIT(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim):\n",
    "        super(AIT, self).__init__()\n",
    "        self.input_size, self.hidden_dim = input_dim, hidden_dim\n",
    "        self.h1 = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.h2 = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.h3 = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.transfer_nn = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))\n",
    "    def forward(self, p, q): \n",
    "        p1, p2, p3 = self.h1(p), self.h2(p), self.h3(p)\n",
    "        q1, q2, q3 = self.h1(q), self.h2(q), self.h3(q)\n",
    "        w_p = torch.sum(p2 * p3, dim=-1, keepdim=True) / np.sqrt(p2.shape[-1])\n",
    "        w_q = torch.sum(q2 * q3, dim=-1, keepdim=True) / np.sqrt(q2.shape[-1])\n",
    "        w_p, w_q = w_p / (w_p + w_q), w_q / (w_p + w_q) # 注意力权重\n",
    "        z = self.transfer_nn(w_p * p1 + w_q * q1)\n",
    "        return z\n",
    "class AITM(Module):\n",
    "    def __init__(self, dense_feature_cols:[(int,int)], sparse_feature_cols:[(int,int)], sparse_feature_embedding_dim, \n",
    "                 hidden_dim:int, num_task:int, dnn_layer_dims:list[int], dnn_dropout=0.):\n",
    "        super(AITM, self).__init__()\n",
    "        self.dense_feature_cols, self.sparse_feature_cols, self.sparse_feature_embedding_dim = dense_feature_cols, sparse_feature_cols, sparse_feature_embedding_dim\n",
    "        self.num_task, self.dnn_layer_dims, self.hidden_dim = num_task, dnn_layer_dims, hidden_dim\n",
    "        # sparse feature embedding dict\n",
    "        self.embed_layers = nn.ModuleDict({'embed_' + str(i): nn.Embedding(num_embeddings=valcount, embedding_dim=sparse_feature_embedding_dim) for i, valcount in sparse_feature_cols})\n",
    "        self.input_dim = len(dense_feature_cols) + len(sparse_feature_cols) * sparse_feature_embedding_dim\n",
    "        # dnn for each task： towers\n",
    "        self.dnn_nets = []\n",
    "        for i in range(num_task):\n",
    "            net = nn.Sequential()\n",
    "            pre_layer_dim = self.input_dim\n",
    "            for layer_dim in dnn_layer_dims:\n",
    "                net.append(nn.Linear(pre_layer_dim, layer_dim))\n",
    "                net.append(nn.BatchNorm1d(layer_dim))\n",
    "                net.append(nn.Dropout(dnn_dropout))\n",
    "                pre_layer_dim = layer_dim\n",
    "            net.append(nn.Linear(dnn_layer_dims[-1], hidden_dim))\n",
    "            net.append(nn.Sigmoid())\n",
    "            self.dnn_nets.append(net)\n",
    "        # AITM\n",
    "        self.aits = [AIT(hidden_dim, hidden_dim) for i in range(num_task)]\n",
    "        # final dnn\n",
    "        self.final_dnns = [nn.Sequential(nn.Linear(hidden_dim, 1), nn.Sigmoid()) for i in range(num_task)]\n",
    "    def forward(self, x):\n",
    "        batch_len = x.shape[0]\n",
    "        dense_input = x[:, :len(self.dense_feature_cols)]\n",
    "        sparse_embeds = torch.cat([self.embed_layers['embed_' + str(i)](x[:, i].long()) for i in range(len(self.dense_feature_cols), x.shape[1])], axis=1)\n",
    "        x = torch.cat([sparse_embeds, dense_input], axis=-1)\n",
    "        towers_out = [dnn_net(x) for dnn_net in self.dnn_nets]\n",
    "        # 第0个,令p = q\n",
    "        outputs = []\n",
    "        for i in range(self.num_task):\n",
    "            if i == 0:\n",
    "                z = self.aits[i](towers_out[i], towers_out[i])\n",
    "            else:\n",
    "                z = self.aits[i](towers_out[i-1], towers_out[i])\n",
    "            outputs.append(self.final_dnns[i](z).squeeze())\n",
    "        return outputs\n",
    "    def parameters(self, recurse: bool = True):\n",
    "        paras = []\n",
    "        for layer in self.embed_layers.values():\n",
    "            paras += [p for p in layer.parameters()]\n",
    "        for net in self.dnn_nets:\n",
    "            paras += [p for p in net.parameters()]\n",
    "        for ait in self.aits:\n",
    "            paras += [p for p in ait.parameters()]\n",
    "        for net in self.final_dnns:\n",
    "            paras += [p for p in net.parameters()]\n",
    "        return paras\n",
    "model = AITM(dense_feature_cols=[i for i in range(num_number_features)], sparse_feature_cols=[(i,len(category_feature_vals[i])) for i in range(num_number_features, num_features)], sparse_feature_embedding_dim =embedding_dim, hidden_dim=embedding_dim, num_task=2, dnn_layer_dims=[128, 32], dnn_dropout=0.).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = MSELoss(reduction='sum').to(device)\n",
    "alpha = 0.5\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    # train:\n",
    "    epoch_train_losses = []\n",
    "    model.train()\n",
    "    for i, inputs in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1, norm_type=2)\n",
    "        optimizer.step()\n",
    "        epoch_train_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    # validate:\n",
    "    model.eval()\n",
    "    epoch_test_losses = []\n",
    "    for i, inputs in enumerate(test_loader):\n",
    "        input = inputs[0].to(device)\n",
    "        label_1, label_2 = inputs[1].to(device), inputs[2].to(device)\n",
    "        output1, output2 = model(input)\n",
    "        loss_1, loss_2 = criterion(output1, label_1), criterion(output2, label_2)\n",
    "        loss = alpha * loss_1 + loss_2\n",
    "        epoch_test_losses.append([input.shape[0], loss_1.item(), loss_2.item(), loss.item()])\n",
    "    num_train, num_test = sum([x[0] for x in epoch_train_losses]), sum([x[0] for x in epoch_test_losses])\n",
    "    train_overall_loss = sum([x[-1] for x in epoch_train_losses])/num_train\n",
    "    train_task1_loss = sum([x[1] for x in epoch_train_losses])/num_train\n",
    "    train_task2_loss = sum([x[2] for x in epoch_train_losses])/num_train\n",
    "    test_overall_loss = sum([x[-1] for x in epoch_test_losses])/num_test\n",
    "    test_task1_loss = sum([x[1] for x in epoch_test_losses])/num_test\n",
    "    test_task2_loss = sum([x[2] for x in epoch_test_losses])/num_test\n",
    "    # print\n",
    "    print('['+datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")+']', 'epoch=[{}/{}], train_mse_overall_loss: {:.6f}, train_mse_task1_loss: {:.6f}, train_mse_task2_loss: {:.6f}, validate_mse_overall_loss: {:.6f}, validate_mse_task1_loss: {:.6f}, validate_mse_task2_loss: {:.6f}'.format(epoch+1, num_epochs,  train_overall_loss, train_task1_loss, train_task2_loss, test_overall_loss, test_task1_loss, test_task2_loss))\n",
    "    gc.collect()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-07T04:27:56.523896800Z",
     "start_time": "2023-09-07T04:24:57.388345400Z"
    }
   },
   "id": "6ee215880c96bd80"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# \n",
    "\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "40d1c5cb82fcf9b2"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}