12011915 周家琛

README.md

@@ -24,7 +24,7 @@ This assignment focuses on implementing and comparing three recommender system m
 
 To run the models, execute the command like:
 ```
-python run_models.py
+python main.py --batch_size=256 --lr=0.001 --factor_num=16
 ```
 
 This script will train the GMF, MLP, and NeuMF models on the MovieLens dataset and output the evaluation metrics HR@10 and NDCG@10 for each model.

config.py

@@ -0,0 +1,18 @@
+# dataset name 
+dataset = 'ml-1m'
+
+# model name 
+model = 'NeuMF-end'
+assert model in ['MLP', 'GMF', 'NeuMF-end']
+
+# paths
+main_path = '/data/lab/NCF_example1/Data/'
+
+train_rating = main_path + '{}.train.rating'.format(dataset)
+test_rating = main_path + '{}.test.rating'.format(dataset)
+test_negative = main_path + '{}.test.negative'.format(dataset)
+
+model_path = './models/'
+GMF_model_path = model_path + 'GMF.pth'
+MLP_model_path = model_path + 'MLP.pth'
+NeuMF_model_path = model_path + 'NeuMF.pth'

data_utils.py

@@ -0,0 +1,84 @@
+import numpy as np 
+import pandas as pd 
+import scipy.sparse as sp
+
+import torch.utils.data as data
+
+import config
+
+
+def load_all(test_num=100):
+    """ We load all the three file here to save time in each epoch. """
+    train_data = pd.read_csv(
+        config.train_rating, 
+        sep='\t', header=None, names=['user', 'item'], 
+        usecols=[0, 1], dtype={0: np.int32, 1: np.int32})
+
+    user_num = train_data['user'].max() + 1
+    item_num = train_data['item'].max() + 1
+
+    train_data = train_data.values.tolist()
+
+    # load ratings as a dok matrix
+    train_mat = sp.dok_matrix((user_num, item_num), dtype=np.float32)
+    for x in train_data:
+        train_mat[x[0], x[1]] = 1.0
+
+    test_data = []
+    with open(config.test_negative, 'r') as fd:
+        line = fd.readline()
+        while line != None and line != '':
+            arr = line.split('\t')
+            u = eval(arr[0])[0]
+            test_data.append([u, eval(arr[0])[1]])
+            for i in arr[1:]:
+                test_data.append([u, int(i)])
+            line = fd.readline()
+    return train_data, test_data, user_num, item_num, train_mat
+
+
+class NCFData(data.Dataset):
+    def __init__(self, features, 
+                num_item, train_mat=None, num_ng=0, is_training=None):
+        super(NCFData, self).__init__()
+        """ Note that the labels are only useful when training, we thus 
+            add them in the ng_sample() function.
+        """
+        self.features_ps = features
+        self.num_item = num_item
+        self.train_mat = train_mat
+        self.num_ng = num_ng
+        self.is_training = is_training
+        self.labels = [0 for _ in range(len(features))]
+
+    def ng_sample(self):
+        assert self.is_training, 'no need to sampling when testing'
+
+        self.features_ng = []
+        for x in self.features_ps:
+            u = x[0]
+            for t in range(self.num_ng):
+                j = np.random.randint(self.num_item)
+                while (u, j) in self.train_mat:
+                    j = np.random.randint(self.num_item)
+                self.features_ng.append([u, j])
+
+        labels_ps = [1 for _ in range(len(self.features_ps))]
+        labels_ng = [0 for _ in range(len(self.features_ng))]
+
+        self.features_fill = self.features_ps + self.features_ng
+        self.labels_fill = labels_ps + labels_ng
+
+    def __len__(self):
+        return (self.num_ng + 1) * len(self.labels)
+
+    def __getitem__(self, idx):
+        features = self.features_fill if self.is_training \
+                    else self.features_ps
+        labels = self.labels_fill if self.is_training \
+                    else self.labels
+
+        user = features[idx][0]
+        item = features[idx][1]
+        label = labels[idx]
+        return user, item ,label

evaluate.py

@@ -0,0 +1,34 @@
+import numpy as np
+import torch
+
+
+def hit(gt_item, pred_items):
+    if gt_item in pred_items:
+        return 1
+    return 0
+
+
+def ndcg(gt_item, pred_items):
+    if gt_item in pred_items:
+        index = pred_items.index(gt_item)
+        return np.reciprocal(np.log2(index+2))
+    return 0
+
+
+def metrics(model, test_loader, top_k):
+    HR, NDCG = [], []
+
+    for user, item, label in test_loader:
+        user = user.cuda()
+        item = item.cuda()
+
+        predictions = model(user, item)
+        _, indices = torch.topk(predictions, top_k)
+        recommends = torch.take(
+                item, indices).cpu().numpy().tolist()
+
+        gt_item = item[0].item()
+        HR.append(hit(gt_item, recommends))
+        NDCG.append(ndcg(gt_item, recommends))
+
+    return np.mean(HR), np.mean(NDCG)

main.py

@@ -0,0 +1,131 @@
+import os
+import time
+import argparse
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data as data
+import torch.backends.cudnn as cudnn
+from tensorboardX import SummaryWriter
+
+import model
+import config
+import evaluate
+import data_utils
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--lr", 
+    type=float, 
+    default=0.001, 
+    help="learning rate")
+parser.add_argument("--dropout", 
+    type=float,
+    default=0.0,  
+    help="dropout rate")
+parser.add_argument("--batch_size", 
+    type=int, 
+    default=256, 
+    help="batch size for training")
+parser.add_argument("--epochs", 
+    type=int,
+    default=20,  
+    help="training epoches")
+parser.add_argument("--top_k", 
+    type=int, 
+    default=10, 
+    help="compute metrics@top_k")
+parser.add_argument("--factor_num", 
+    type=int,
+    default=32, 
+    help="predictive factors numbers in the model")
+parser.add_argument("--num_layers", 
+    type=int,
+    default=3, 
+    help="number of layers in MLP model")
+parser.add_argument("--num_ng", 
+    type=int,
+    default=4, 
+    help="sample negative items for training")
+parser.add_argument("--test_num_ng", 
+    type=int,
+    default=99, 
+    help="sample part of negative items for testing")
+parser.add_argument("--out", 
+    default=True,
+    help="save model or not")
+parser.add_argument("--gpu", 
+    type=str,
+    default="0",  
+    help="gpu card ID")
+args = parser.parse_args()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+cudnn.benchmark = True
+
+
+############################## PREPARE DATASET ##########################
+train_data, test_data, user_num ,item_num, train_mat = data_utils.load_all()
+
+# construct the train and test datasets
+train_dataset = data_utils.NCFData(
+        train_data, item_num, train_mat, args.num_ng, True)
+test_dataset = data_utils.NCFData(
+        test_data, item_num, train_mat, 0, False)
+train_loader = data.DataLoader(train_dataset,
+        batch_size=args.batch_size, shuffle=True, num_workers=4)
+test_loader = data.DataLoader(test_dataset,
+        batch_size=args.test_num_ng+1, shuffle=False, num_workers=0)
+
+########################### CREATE MODEL #################################
+GMF_model = None
+MLP_model = None
+
+model = model.NCF(user_num, item_num, args.factor_num, args.num_layers, 
+                        args.dropout, config.model, GMF_model, MLP_model)
+model.cuda()
+loss_function = nn.BCEWithLogitsLoss()
+
+optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+writer = SummaryWriter() # for visualization
+
+########################### TRAINING #####################################
+count, best_hr = 0, 0
+for epoch in range(args.epochs):
+    model.train() # Enable dropout (if have).
+    start_time = time.time()
+    train_loader.dataset.ng_sample()
+
+    for user, item, label in train_loader:
+        user = user.cuda()
+        item = item.cuda()
+        label = label.float().cuda()
+
+        model.zero_grad()
+        prediction = model(user, item)
+        loss = loss_function(prediction, label)
+        loss.backward()
+        optimizer.step()
+        writer.add_scalar('data/loss', loss.item(), count)
+        count += 1
+
+    model.eval()
+    HR, NDCG = evaluate.metrics(model, test_loader, args.top_k)
+
+    elapsed_time = time.time() - start_time
+    print("The time elapse of epoch {:03d}".format(epoch) + " is: " + 
+            time.strftime("%H: %M: %S", time.gmtime(elapsed_time)))
+    print("HR: {:.3f}\tNDCG: {:.3f}".format(np.mean(HR), np.mean(NDCG)))
+
+    if HR > best_hr:
+        best_hr, best_ndcg, best_epoch = HR, NDCG, epoch
+        if args.out:
+            if not os.path.exists(config.model_path):
+                os.mkdir(config.model_path)
+            torch.save(model, 
+                '{}{}.pth'.format(config.model_path, config.model))
+
+print("End. Best epoch {:03d}: HR = {:.3f}, NDCG = {:.3f}".format(best_epoch, best_hr, best_ndcg))

model.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F 
+
+
+class NCF(nn.Module):
+    def __init__(self, user_num, item_num, factor_num, num_layers,
+                    dropout, model, GMF_model=None, MLP_model=None):
+        super(NCF, self).__init__()
+        """
+        user_num: number of users;
+        item_num: number of items;
+        factor_num: number of predictive factors;
+        num_layers: the number of layers in MLP model;
+        dropout: dropout rate between fully connected layers;
+        model: 'MLP', 'GMF', 'NeuMF-end';
+        """
+        self.dropout = dropout
+        self.model = model
+        self.GMF_model = GMF_model
+        self.MLP_model = MLP_model
+
+        self.embed_user_GMF = nn.Embedding(user_num, factor_num)
+        self.embed_item_GMF = nn.Embedding(item_num, factor_num)
+        self.embed_user_MLP = nn.Embedding(
+                user_num, factor_num * (2 ** (num_layers - 1)))
+        self.embed_item_MLP = nn.Embedding(
+                item_num, factor_num * (2 ** (num_layers - 1)))
+
+        MLP_modules = []
+        for i in range(num_layers):
+            input_size = factor_num * (2 ** (num_layers - i))
+            MLP_modules.append(nn.Dropout(p=self.dropout))
+            MLP_modules.append(nn.Linear(input_size, input_size//2))
+            MLP_modules.append(nn.ReLU())
+        self.MLP_layers = nn.Sequential(*MLP_modules)
+
+        if self.model in ['MLP', 'GMF']:
+            predict_size = factor_num 
+        else:
+            predict_size = factor_num * 2
+        self.predict_layer = nn.Linear(predict_size, 1)
+
+        self._init_weight_()
+
+    def _init_weight_(self):
+        """ We leave the weights initialization here. """
+        nn.init.normal_(self.embed_user_GMF.weight, std=0.01)
+        nn.init.normal_(self.embed_user_MLP.weight, std=0.01)
+        nn.init.normal_(self.embed_item_GMF.weight, std=0.01)
+        nn.init.normal_(self.embed_item_MLP.weight, std=0.01)
+
+        for m in self.MLP_layers:
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+
+        nn.init.kaiming_uniform_(self.predict_layer.weight, a=1, nonlinearity='sigmoid')
+
+        for m in self.modules():
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                m.bias.data.zero_()
+
+
+    def forward(self, user, item):
+        if not self.model == 'MLP':
+            embed_user_GMF = self.embed_user_GMF(user)
+            embed_item_GMF = self.embed_item_GMF(item)
+            output_GMF = embed_user_GMF * embed_item_GMF
+        if not self.model == 'GMF':
+            embed_user_MLP = self.embed_user_MLP(user)
+            embed_item_MLP = self.embed_item_MLP(item)
+            interaction = torch.cat((embed_user_MLP, embed_item_MLP), -1)
+            output_MLP = self.MLP_layers(interaction)
+
+        if self.model == 'GMF':
+            concat = output_GMF
+        elif self.model == 'MLP':
+            concat = output_MLP
+        else:
+            concat = torch.cat((output_GMF, output_MLP), -1)
+
+        prediction = self.predict_layer(concat)
+        return prediction.view(-1)