train_list.py

import os
import numpy as np
import random

import torch
import torch.nn as nn

from model import *

import arguments

import utils.load_dataset
import utils.data_loader
import utils.metrics
from utils.early_stop import EarlyStopping, Stop_args

def setup_seed(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available(): 
        torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

def para(args): 
    if args.dataset == 'simulation': 
        args.training_args = {'batch_size': 1024, 'epochs': 500, 'patience': 60, 'block_batch': [20, 500]}
        args.base_model_args = {'emb_dim': 10, 'learning_rate': 0.01, 'imputaion_lambda': 0.01, 'weight_decay': 1}
        args.weight1_model_args = {'learning_rate': 0.1, 'weight_decay': 0.001}
        args.weight2_model_args = {'learning_rate': 1e-3, 'weight_decay': 1e-2}
        args.imputation_model_args = {'learning_rate': 1e-1, 'weight_decay': 1e-4}
    else: 
        print('invalid arguments')
        os._exit()

def train_and_eval(train_data, unif_train_data, val_data, test_data, device = 'cuda', 
        base_model_args: dict = {'emb_dim': 64, 'dropout': 0.0, 'learning_rate': 0.05, 'imputaion_lambda': 0.01, 'weight_decay': 0.05}, 
        weight1_model_args: dict = {'learning_rate': 0.1, 'weight_decay': 0.005}, 
        weight2_model_args: dict = {'learning_rate': 0.1, 'weight_decay': 0.005}, 
        imputation_model_args: dict = {'learning_rate': 0.01, 'weight_decay': 0.5}, 
        training_args: dict =  {'batch_size': 1024, 'epochs': 100, 'patience': 20, 'block_batch': [1000, 100]}): 
    train_position = train_data['position']
    train_rating = train_data['rating']
    train_dense = train_rating.to_dense()
    # uniform data
    users_unif = unif_train_data._indices()[0]
    items_unif = unif_train_data._indices()[1]
    y_unif = unif_train_data._values()

    # build data_loader. 
    train_loader = utils.data_loader.User(train_position, train_rating, u_batch_size=training_args['block_batch'][0], device=device)
    val_loader = utils.data_loader.DataLoader(utils.data_loader.Interactions(val_data), batch_size=training_args['batch_size'], shuffle=False, num_workers=0)
    test_loader = utils.data_loader.DataLoader(utils.data_loader.Interactions(test_data), batch_size=training_args['batch_size'], shuffle=False, num_workers=0)

    n_user, n_item = train_position.shape
    n_position = torch.max(train_position._values()).item() + 1

    # base model and its optimizer
    base_model = MetaMF(n_user, n_item, dim=base_model_args['emb_dim'], dropout=0).to(device)
    base_optimizer = torch.optim.SGD(base_model.params(), lr=base_model_args['learning_rate'], weight_decay=0) # todo: other optimizer SGD

    # meta models and their optimizer
    weight1_model = FourLinear(n_user, n_item, 2, n_position = n_position).to(device)
    weight1_optimizer = torch.optim.Adam(weight1_model.parameters(), lr=weight1_model_args['learning_rate'], weight_decay=weight1_model_args['weight_decay'])

    weight2_model = ThreeLinear(n_user, n_item, 2).to(device)
    weight2_optimizer = torch.optim.Adam(weight2_model.parameters(), lr=weight2_model_args['learning_rate'], weight_decay=weight2_model_args['weight_decay'])

    imputation_model = OneLinear(3).to(device)
    imputation_optimizer = torch.optim.Adam(imputation_model.parameters(), lr=imputation_model_args['learning_rate'], weight_decay=imputation_model_args['weight_decay'])
    
    # loss_criterion
    sum_criterion = nn.MSELoss(reduction='sum')
    none_criterion = nn.MSELoss(reduction='none')

    # begin training
    stopping_args = Stop_args(patience=training_args['patience'], max_epochs=training_args['epochs'])
    early_stopping = EarlyStopping(base_model, **stopping_args)

    for epo in range(early_stopping.max_epochs):
        training_loss = 0 
        for u_batch_idx, users in enumerate(train_loader.User_loader): 
            users_train, items_train, positions_train, y_train = train_loader.get_batch(users)

            # all pair
            all_pair = torch.cartesian_prod(users, torch.arange(n_item).to(device))
            users_all, items_all = all_pair[:,0], all_pair[:,1]
            
            # calculate weight1
            weight1_model.train()
            weight1 = weight1_model(users_train, items_train, (y_train==1)*1, positions_train)
            weight1 = torch.exp(weight1/5) # for stable training

            # calculate weight2
            weight2_model.train()
            weight2 = weight2_model(users_all, items_all,(train_dense[users_all,items_all]!=0)*1)
            weight2 = torch.exp(weight2/5) #for stable training

            # calculate imputation values
            imputation_model.train()
            impu_f_all = torch.tanh(imputation_model((train_dense[users_all,items_all]).long()+1))

            # one_step_model: assumed model, just update one step on base model. it is for updating weight parameters
            one_step_model = MetaMF(n_user, n_item, dim=base_model_args['emb_dim'], dropout=0)
            one_step_model.load_state_dict(base_model.state_dict())

            # formal parameter: Using training set to update parameters
            one_step_model.train()
            # all pair data in this block
            y_hat_f_all = one_step_model(users_all, items_all)
            cost_f_all = none_criterion(y_hat_f_all, impu_f_all)
            loss_f_all = torch.sum(cost_f_all * weight2)
            # observation data
            y_hat_f_obs = one_step_model(users_train, items_train)
            cost_f_obs = none_criterion(y_hat_f_obs, y_train)
            loss_f_obs = torch.sum(cost_f_obs * weight1)
            loss_f = loss_f_obs + base_model_args['imputaion_lambda'] * loss_f_all + base_model_args['weight_decay'] * one_step_model.l2_norm(users_all, items_all)
            
            # update parameters of one_step_model
            one_step_model.zero_grad()
            grads = torch.autograd.grad(loss_f, (one_step_model.params()), create_graph=True)
            one_step_model.update_params(base_model_args['learning_rate'], source_params=grads)

            # latter hyper_parameter: Using uniform set to update hyper_parameters
            y_hat_l = one_step_model(users_unif, items_unif)
            loss_l = sum_criterion(y_hat_l, y_unif)
            
            # update weight parameters
            weight1_optimizer.zero_grad()
            weight2_optimizer.zero_grad()
            imputation_optimizer.zero_grad()
            loss_l.backward()
            
            if epo>=20:
                weight1_optimizer.step()
                weight2_optimizer.step()
            imputation_optimizer.step()

            # use new weights to update parameters        
            weight1_model.train()
            weight1 = weight1_model(users_train, items_train, (y_train==1)*1, positions_train)
            weight1 = torch.exp(weight1/5) # for stable training
            
            # calculate weight2
            weight2_model.train()
            weight2 = weight2_model(users_all, items_all,(train_dense[users_all,items_all]!=0)*1)
            weight2 = torch.exp(weight2/5) # for stable training
            
            # use new imputation to update parameters
            imputation_model.train()
            impu_all = torch.tanh(imputation_model((train_dense[users_all,items_all]).long()+1))

            # loss of training set
            base_model.train()
            # all pair
            y_hat_all = base_model(users_all, items_all)
            cost_all = none_criterion(y_hat_all, impu_all)
            loss_all = torch.sum(cost_all * weight2)
            # observation
            y_hat_obs = base_model(users_train, items_train)
            cost_obs = none_criterion(y_hat_obs, y_train)
            loss_obs = torch.sum(cost_obs * weight1)
            loss = loss_obs + base_model_args['imputaion_lambda'] * loss_all + base_model_args['weight_decay'] * base_model.l2_norm(users_all, items_all)
            
            base_optimizer.zero_grad()
            loss.backward()
            base_optimizer.step()

            training_loss += loss.item()            

        
        base_model.eval()
        with torch.no_grad():
            # training metrics
            train_pre_ratings = torch.empty(0).to(device)
            train_ratings = torch.empty(0).to(device)
            for u_batch_idx, users in enumerate(train_loader.User_loader): 
                users_train, items_train, positions_train, y_train = train_loader.get_batch(users)
                pre_ratings = base_model(users_train, items_train)
                train_pre_ratings = torch.cat((train_pre_ratings, pre_ratings))
                train_ratings = torch.cat((train_ratings, y_train))

            # validation metrics
            val_pre_ratings = torch.empty(0).to(device)
            val_ratings = torch.empty(0).to(device)
            for batch_idx, (users, items, ratings) in enumerate(val_loader):
                pre_ratings = base_model(users, items)
                val_pre_ratings = torch.cat((val_pre_ratings, pre_ratings))
                val_ratings = torch.cat((val_ratings, ratings))
            
        train_results = utils.metrics.evaluate(train_pre_ratings, train_ratings, ['MSE', 'NLL'])
        val_results = utils.metrics.evaluate(val_pre_ratings, val_ratings, ['MSE', 'NLL', 'AUC'])

        print('Epoch: {0:2d} / {1}, Traning: {2}, Validation: {3}'.
                format(epo, training_args['epochs'], ' '.join([key+':'+'%.3f'%train_results[key] for key in train_results]), 
                ' '.join([key+':'+'%.3f'%val_results[key] for key in val_results])))

        if epo >= 50 and early_stopping.check([val_results['AUC']], epo):
            break


    # restore best model
    print('Loading {}th epoch'.format(early_stopping.best_epoch))
    base_model.load_state_dict(early_stopping.best_state)

    # validation metrics
    val_pre_ratings = torch.empty(0).to(device)
    val_ratings = torch.empty(0).to(device)
    for batch_idx, (users, items, ratings) in enumerate(val_loader):
        pre_ratings = base_model(users, items)
        val_pre_ratings = torch.cat((val_pre_ratings, pre_ratings))
        val_ratings = torch.cat((val_ratings, ratings))

    # test metrics
    test_users = torch.empty(0, dtype=torch.int64).to(device)
    test_items = torch.empty(0, dtype=torch.int64).to(device)
    test_pre_ratings = torch.empty(0).to(device)
    test_ratings = torch.empty(0).to(device)
    for batch_idx, (users, items, ratings) in enumerate(test_loader):
        pre_ratings = base_model(users, items)
        test_users = torch.cat((test_users, users))
        test_items = torch.cat((test_items, items))
        test_pre_ratings = torch.cat((test_pre_ratings, pre_ratings))
        test_ratings = torch.cat((test_ratings, ratings))

    val_results = utils.metrics.evaluate(val_pre_ratings, val_ratings, ['MSE', 'NLL', 'AUC'])
    test_results = utils.metrics.evaluate(test_pre_ratings, test_ratings, ['MSE', 'NLL', 'AUC', 'Recall_Precision_NDCG@'], users=test_users, items=test_items)
    print('-'*30)
    print('The performance of validation set: {}'.format(' '.join([key+':'+'%.3f'%val_results[key] for key in val_results])))
    print('The performance of testing set: {}'.format(' '.join([key+':'+'%.3f'%test_results[key] for key in test_results])))
    print('-'*30)
    return val_results,test_results

if __name__ == "__main__": 
    args = arguments.parse_args()
    para(args)
    setup_seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train, unif_train, validation, test = utils.load_dataset.load_dataset(data_name=args.dataset, type = 'list', seed = args.seed, device=device)
    train_and_eval(train, unif_train, validation, test, device, base_model_args = args.base_model_args, 
                weight1_model_args = args.weight1_model_args, weight2_model_args = args.weight2_model_args, imputation_model_args = args.imputation_model_args, training_args = args.training_args)