main.py

import logging
import datetime
import os
from tqdm import tqdm
import json
import math
from json import dumps as json_dumps
import torch
import torch.optim as optim
from parses import parse_args
from dataprocess import DataProcesser
from model import D2ASR, GaussianDiffusion, Denoise, Extractor
from utils import *
from evaluation import calcRes
import warnings
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import optuna
best_recall = None
warnings.filterwarnings("ignore", category=RuntimeWarning)

args = parse_args()
arrange_seed(args.seed)

class Trainer:
    def __init__(self, args, config, data_processer) -> None:
        ## Set Arguments
        self.args = args
        
        ## Single valued parameter
        self.device = args.device
        self.topk = args.topk
        self.batch_size = args.batch_size
        self.test_batch_size = args.test_batch_size
        self.user_num = config['user_num']
        self.item_num = config['item_num']
        
        ## Data Processer
        self.train_loader = data_processer.trainLoader
        self.test_loader = data_processer.testLoader
        self.diff_loader = data_processer.diffusionLoader
        
        ## Coo Adjusted Graphs
        self.adj_graph = config['adj_graph']
        self.social_graph = config['social_graph']
        self.social_dict = config['social_dict']
        
        ## Torch Sparse Graphs
        self.ui_matrix = data_processer.ui_matrix
        self.norm_adj = data_processer.norm_adj
        self.norm_social = data_processer.norm_social
        in_degree = data_processer.in_degree
        in_degree_social = data_processer.social_in_degree
        self.pop_emb_inter = self.pop_enc(in_degree, args.embed_dim)
        self.pop_emb_social = self.pop_enc(in_degree_social, args.embed_dim)
        
        ## Model and Related Optimization
        
        self.model = D2ASR(config, args).to(self.device)
        self.aug_1 = Extractor([64, 32, 16, 64]).to(self.device)
        self.aug_2 = Extractor([64, 32, 16, 64]).to(self.device)
        #self.optimizer = optim.Adam(self.model.parameters(), lr=args.lr)
        self.optimizer = optim.Adam(self.model.parameters(), lr=args.lr, weight_decay=0)
        self.optimizer_1 = optim.Adam([{"params": self.aug_1.parameters()}, {"params": self.aug_2.parameters()}], lr=args.lr)
        self.optimizer_2 = optim.Adam([{"params": self.model.parameters()}, {"params": self.aug_1.parameters()}, {"params": self.aug_2.parameters()}], lr=args.lr)
        self.diffusion_model = GaussianDiffusion(self.device, self.item_num, args.noise_scale, args.noise_min, args.noise_max, args.steps).to(self.device)
        out_dims = eval(args.dims) + [args.embed_dim]
        in_dims = out_dims[::-1]
        z_dim = args.z_dim
        self.denoise_model = Denoise(self.device, args.n_layer_sn, args.dropout_rate, config['social_graph'], in_dims, out_dims, z_dim, args.diff_emb_size, norm=args.norm).to(self.device)
        self.denoise_opt = optim.Adam(self.denoise_model.parameters(), lr=args.lr_diff, weight_decay=0)
        
        print('Done with device: {}'.format(self.device))
        
        ## Drawing & Painting
        self.loss4epoch_dict = {
            "Rec" : [],
            "Pre" : [],
            "BPR" : [],
            "Diff" : [],
            "Us" : [],
            "CL" : [],
            "KD" : []
        }
        self.loss4epoch_path = "./draw/loss4epoch.json"

    
    def save_loss(self, filepath=None):
        if filepath == None:
            with open(self.loss4epoch_path, 'w', encoding="utf-8")as f:
                json.dump(self.loss4epoch_dict, f, indent=4, ensure_ascii=False)
        else:
            with open(filepath, 'w', encoding="utf-8")as f:
                json.dump(self.loss4epoch_dict, f, indent=4, ensure_ascii=False)
    
    def save_A(self, A, epoch, dataset):
        # 将参数A保存到CSV文件中
        A_numpy = A.detach().cpu().numpy()
        df = pd.DataFrame(A_numpy)
        df.to_csv(f'./draw/csv/{dataset}/A_epoch_{epoch}.csv', index=False)
        
        # 生成热力图
        plt.figure(figsize=(10, 8))  # 调整图像大小
        plt.imshow(A_numpy, cmap='viridis', aspect='auto')
        plt.colorbar()
        plt.title(f'Heatmap of A at Epoch {epoch}')
        plt.savefig(f'./draw/fig/{dataset}/A_epoch_{epoch}.png')
        plt.close()
    
    def sampSocialGraph(self, uuMat, batch_size):
        """Sampling on social graph for calculate BPRLoss and self-augmented learning loss.
        
        Args:
            uuMat: An user-user matrix.

        Returns:
            Users and user pairs (edges) sampled on social graph.
        """
        batIdx = torch.randint(high=len(uuMat.data), size=(batch_size, ))
        usr0 = torch.from_numpy(uuMat.row[batIdx])
        usrP = torch.from_numpy(uuMat.col[batIdx])
        usrN = torch.from_numpy(self.train_loader.dataset.uuNegs[batIdx])
        usr1 = torch.randint(high=uuMat.shape[0], size=(batch_size,))
        usr2 = torch.randint(high=uuMat.shape[1], size=(batch_size,))
        return usr0, usrP, usrN, usr1, usr2
    
    def pop_enc(self, in_degree, d_model):
        """
        Input: batch*seq_len.
        Output: batch*seq_len*d_model.
        """
        pop_vec = torch.tensor(
            [math.pow(10000.0, 2.0 * (i // 2) / d_model) for i in range(d_model)],
            device=self.device, dtype=torch.float16)

        result = in_degree.unsqueeze(-1) / pop_vec
        result[:, 0::2] = torch.sin(result[:, 0::2])
        result[:, 1::2] = torch.cos(result[:, 1::2])
        return result
    
    def train_single_epoch(self, dataloader, diffloader):
        self.model.train()
        
        # Initialize the best loss, mf_loss, emb_loss, cen_loss, kd_loss
        loss, rec_loss, cl_loss, diff_loss, us_loss = 0., 0., 0., 0., 0.
        
        # Initialize the negative label sampling
        dataloader.dataset.neg_sampling()
        steps = dataloader.dataset.__len__() // args.batch_size
        dropped_adj1 = self.model.edge_dropout(self.adj_graph)
        dropped_adj2 = self.model.edge_dropout(self.adj_graph)
        dropped_social1 = self.model.edge_dropout(self.social_graph, social=True)
        dropped_social2 = self.model.edge_dropout(self.social_graph, social=True)
        
        ''' Start batch training - Diffusion'''
        if args.use_diff:
            #pbar = tqdm(diffloader, desc='Diff-Epoch #{}, processing'.format(self.epoch))
            pbar = tqdm(dataloader, desc='Diff-Epoch #{}, processing'.format(self.epoch))
            for i, batch in enumerate(pbar):
                #batch_user, batch_index = batch
                #batch_user, batch_index = batch_user.to(self.device), batch_index.to(self.device)
                ui_user_idx, ui_pos_idx, ui_neg_idx= batch
                
                ui_matrix = self.ui_matrix
                item_embedding = self.model.get_item_embedding()
                user_embedding = self.model.get_user_embedding()
                user_social_embeding = self.denoise_model.get_social_embedding(user_embedding)
                batch_user_emb = user_social_embeding[ui_user_idx]
                
                self.denoise_opt.zero_grad()
                batch_diff_loss, batch_us_loss = self.diffusion_model.training_losses(self.denoise_model, batch_user_emb, ui_matrix, user_embedding, item_embedding, ui_user_idx)
                batch_loss = batch_diff_loss * (1-args.e_loss) + batch_us_loss * args.e_loss
                
                diff_loss += batch_diff_loss.item()
                us_loss += batch_us_loss.item()
                
                pbar.set_description("Diff-Epoch #{}, batch_loss {:.4f}, diff_loss {:.4f}, us_loss {:.4f}".format(self.epoch, batch_loss.item(), batch_diff_loss.item(), batch_us_loss.item()))
                batch_loss.backward()
                self.denoise_opt.step()
                self.loss4epoch_dict["Pre"].append(batch_loss.item())
                self.loss4epoch_dict["Diff"].append(batch_diff_loss.item())
                self.loss4epoch_dict["Us"].append(batch_us_loss.item())
                
            ''' Rebuild Social Network'''
            with torch.no_grad():
                #denoised_edges = []
                #h_list = []
                #t_list = []
                user_embedding = self.model.get_user_embedding()
                user_social_embeding = self.denoise_model.get_social_embedding(user_embedding)
                denoised_social_norm = self.diffusion_model.p_sample(self.denoise_model, user_social_embeding, args.sampling_steps)
                #top_user, indices_ = torch.topk(denoised_batch, k=args.rebuild_k)
                #for i in range(batch_index.shape[0]):
                #    for j in range(indices_[i].shape[0]):
                #        h_list.append(batch_index[i])
                #        t_list.append(indices_[i][j])

                """ edge_set = set()
                for index in range(len(h_list)):
                    edge_set.add((int(h_list[index].cpu().numpy()), int(t_list[index].cpu().numpy())))
                for index in range(len(h_list)):
                    if (int(t_list[index].cpu().numpy()), int(h_list[index].cpu().numpy())) not in edge_set:
                        h_list.append(t_list[index])
                        t_list.append(h_list[index])
                social_dict = self.social_dict
                for index in range(len(h_list)):
                    try:
                        assert social_dict[int(h_list[index].cpu().numpy())][int(t_list[index].cpu().numpy())]
                        denoised_edges.append([h_list[index], t_list[index]])
                    except Exception:
                        continue
                graph_tensor = torch.tensor(denoised_edges)
                denoised_socialnet = graph_tensor.t().long().to(self.device)
                denoised_social_norm = self.model.plain_to_norm_social(denoised_socialnet) """
            
        
        pbar = tqdm(dataloader, desc='Rec-Epoch  #{}, Processing'.format(self.epoch))
        for i, batch in enumerate(pbar):
            self.batch_id = i
            ui_user_idx, ui_pos_idx, ui_neg_idx= batch
            uu_user_idx, pos_user_idx, neg_user_idx, user_1, user_2 = self.sampSocialGraph(self.social_graph, self.batch_size)
            # Run and Get Loss
            
            ''' == Share Params == '''
            self.optimizer_2.zero_grad()
            if args.use_diff:
                batch_bpr_loss, batch_reg_loss,         \
                batch_social_loss, batch_social_cl_loss, \
                batch_cl_loss, batch_cross_cl_loss,       \
                batch_kd_loss                              = self.model(ui_user_idx, ui_pos_idx, ui_neg_idx, uu_user_idx, pos_user_idx, neg_user_idx, self.pop_emb_inter, self.pop_emb_social, self.norm_adj, dropped_adj1, dropped_adj2, denoised_social_norm, dropped_social2, self.aug_1, self.aug_2, 'step1')
            else:
                batch_bpr_loss, batch_reg_loss,         \
                batch_social_loss, batch_social_cl_loss, \
                batch_cl_loss, batch_cross_cl_loss,       \
                batch_kd_loss                              = self.model(ui_user_idx, ui_pos_idx, ui_neg_idx, uu_user_idx, pos_user_idx, neg_user_idx, self.pop_emb_inter, self.pop_emb_social, self.norm_adj, dropped_adj1, dropped_adj2, dropped_social1, dropped_social2, self.aug_1, self.aug_2, 'step1')
            batch_loss = batch_bpr_loss + batch_social_loss + batch_reg_loss + batch_cl_loss + batch_social_cl_loss + batch_kd_loss + batch_cross_cl_loss
            
            pbar.set_description("Rec-Epoch  #{}, batch_loss {:.4f}, rec_loss {:.4f},  social_loss {:.4f},  cl_loss {:.4f},  social_cl {:.4f},  kd_loss {:.4f}".format(self.epoch, batch_loss.item(), batch_bpr_loss.item(), batch_social_loss.item(), batch_cl_loss.item(), batch_social_cl_loss.item(), batch_kd_loss.item()))

            batch_loss.backward()
            self.optimizer_2.step()
            
            
            self.loss4epoch_dict["Rec"].append(batch_loss.item())
            self.loss4epoch_dict["BPR"].append(batch_bpr_loss.item())
            self.loss4epoch_dict["CL"].append(batch_cl_loss.item())
            self.loss4epoch_dict["KD"].append(batch_kd_loss.item())
                
        result = dict()
        result['loss'] = loss
        result['rec_loss'] = rec_loss
        result['cl_loss'] = cl_loss
        result['diff_loss'] = diff_loss
        result['us_loss'] = us_loss
        #self.save_A(self.denoise_model.A, self.epoch, args.dataset)
        return result
            
    
    def test_epoch(self, testloader):
        epRecall, epNdcg = [0] * 2
        i = 0
        num = testloader.dataset.__len__()
        steps = num // args.test_batch_size

        with torch.no_grad():
            user_embedding, item_embedding = self.model.predict(self.norm_adj)
        pbar = tqdm(testloader, desc='Eval       #{}, top {}'.format(self.epoch, self.topk))
        for user, train_mask in pbar:
            i += 1
            user = user.long().to(self.device)
            train_mask = train_mask.to(self.device)

            allPreds = torch.mm(user_embedding[user], torch.transpose(item_embedding, 1, 0)) * (1 - train_mask) - train_mask * 1e8
            _, topLocs = torch.topk(allPreds, args.topk)
            recall, ndcg = calcRes(topLocs.cpu().numpy(), self.test_loader.dataset.testLocs, user, self.topk)
            epRecall += recall
            epNdcg += ndcg
            pbar.set_description('Eval       #{}, Top {}: Recall {:.4f}, NDCG {:.4f}'.format(self.epoch, self.topk, recall, ndcg))
        result = dict()
        result['Recall'] = epRecall / num
        result['NDCG'] = epNdcg / num
        return result
            
    def train(self, logger):
        args = self.args
        
        ## Initialize the best Evaluation Metrics
        best_hr, best_recall, best_ndcg, best_epoch, wait = 0., 0., 0., 0, 0
        
        ## Initialize Animation
        
        ## Let's Rock Together!
        for self.epoch in range(1, args.n_epoch + 1):
            ## train
            t1 = datetime.datetime.now()
            epoch_losses = self.train_single_epoch(self.train_loader, self.diff_loader)
            t2 = datetime.datetime.now()
            if self.epoch % args.test_epoch == 0 or self.epoch == 1:
                result = self.test_epoch(self.test_loader)
                if result['Recall'] > best_recall:
                    best_recall = result['Recall']
                    best_ndcg = result['NDCG']
                    best_epoch = self.epoch
                print(result)
                # logger.info("Eval Epoch #{}: {}".format(self.epoch, result))
        self.save_loss(None)
        best_msg = "Best Eval: Epoch #{}\tRecall {}\tNDCG {}".format(best_epoch, best_recall, best_ndcg)
        print(best_msg)
        logger.info(best_msg)
        
        
        
        
        
        
    

def objective(trial):
    
    """ args.cross_cl_reg = trial.suggest_float('cross_cl_reg', 1e-7, 1e-2, log=True)
    args.ssl_reg = trial.suggest_float('ssl_reg', 1e-5, 0.5, log=True) """
    args.kd_reg_T2S = trial.suggest_float('kd_reg_T2S', 0.0, 2.0)
    args.kd_reg_S2T = trial.suggest_float('kd_reg_S2T', 0.0, 2.0)
    
    logger = logging.getLogger('train_logger') 
    """ ********* Prepare the log file ********* """
    curr_time = datetime.datetime.now()
    if not os.path.exists('log'):
        os.mkdir('log')
    
    logger.setLevel(logging.INFO)
    logfile = logging.FileHandler('log/{}.log'.format(args.dataset), 'a', encoding='utf-8')
    logfile.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(message)s')
    logfile.setFormatter(formatter)
    logger.addHandler(logfile)

    """ ********* Prepare the dataset ********* """
    print('-' * 20, "Data Information", '-' * 20)
    data_processer = DataProcesser(args)
    logger.info(data_processer.get_summary())
    
    print('-' * 20, "Model Settings", '-' * 20)
    print('{}'.format(
            json_dumps(args.__dict__, indent=4)), flush=True)
    logger.info(args)
    
    """ ********* Generate the normalized adj matrix & social network matrix ********* """
    config = dict()
    config['user_num'] = data_processer.n_users
    config['item_num'] = data_processer.n_items
    plain_adj = data_processer.get_plain_adj_mat()
    social_net = data_processer.get_social_mat()
    social_dict = data_processer.get_social_dict()
    config['adj_graph'] = plain_adj
    config['social_graph'] = social_net
    config['social_dict'] = social_dict
    
    """ ********* Prepare the model and Start training ********* """
    
    print('-' * 20, "Preparing Model", '-' * 20)
    trainer = Trainer(args, config, data_processer)
    print('-' * 20, "Start Training", '-' * 20)
    
    best_recall = trainer.train(logger)

    return best_recall
        
if __name__ == '__main__':

    # 创建 Optuna study 并优化目标函数
    study = optuna.create_study(direction='maximize')
    study.optimize(objective, n_trials=args.n_trials)

    # 获取最佳超参数
    best_params = study.best_params
    print(best_params)
    # print(best_recall)