main.py

# -*- coding: utf-8 -*-
# @Time    : 2021-04-19 17:25
# @Author  : WenYi
# @Contact : 1244058349@qq.com
# @Description :  script description


from utils import data_preparation, TrainDataSet
from torch.utils.data import DataLoader
from model_train import train_model
from esmm import ESMM
from mmoe import MMOE
import torch
import torch.nn as nn


def main():
    train_data, test_data, user_feature_dict, item_feature_dict = data_preparation()
    train_dataset = (train_data.iloc[:, :-2].values, train_data.iloc[:, -2].values, train_data.iloc[:, -1].values)
    # val_dataset = (val_data.iloc[:, :-2].values, val_data.iloc[:, -2].values, val_data.iloc[:, -1].values)
    test_dataset = (test_data.iloc[:, :-2].values, test_data.iloc[:, -2].values, test_data.iloc[:, -1].values)
    train_dataset = TrainDataSet(train_dataset)
    # val_dataset = TrainDataSet(val_dataset)
    test_dataset = TrainDataSet(test_dataset)

    # dataloader
    train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    # val_dataloader = DataLoader(val_dataset, batch_size=64, shuffle=False)
    test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)
    
    # pytorch优化参数
    learn_rate = 0.01
    bce_loss = nn.BCEWithLogitsLoss()
    early_stop = 3
    
    # train model
    # esmm Epoch 17 val loss is 1.164, income auc is 0.875 and marry auc is 0.953
    esmm = ESMM(user_feature_dict, item_feature_dict, emb_dim=64)
    optimizer = torch.optim.Adam(esmm.parameters(), lr=learn_rate)
    train_model(esmm, train_dataloader, test_dataloader, 20, bce_loss, optimizer, 'model/model_esmm_{}', early_stop)
    
    # mmoe
    mmoe = MMOE(user_feature_dict, item_feature_dict, emb_dim=64)
    optimizer = torch.optim.Adam(mmoe.parameters(), lr=learn_rate)
    train_model(mmoe, train_dataloader, test_dataloader, 20, bce_loss, optimizer, 'model/model_mmoe_{}', early_stop)
    

if __name__ == "__main__":
    main()