├── README.md ├── config.py ├── data_utils.py ├── evaluate.py ├── main.py └── model.py /README.md: -------------------------------------------------------------------------------- 1 | # Pytorch-BPR 2 | 3 | Note that I use the two sub datasets provided by Xiangnan's [repo](https://github.com/hexiangnan/neural_collaborative_filtering/tree/master/Data). Another pytorch NCF implementaion can be found at this [repo](https://github.com/guoyang9/NCF). 4 | 5 | I utilized a factor number **32**, and posted the results in the NCF paper and this implementation here. Since there is no specific numbers in their paper, I found this implementation achieved a better performance than the original curve. Moreover, the batch_size is not very sensitive with the final model performance. 6 | 7 | Models | MovieLens HR@10 | MovieLens NDCG@10 | Pinterest HR@10 | Pinterest NDCG@10 8 | ------ | --------------- | ----------------- | --------------- | ----------------- 9 | pytorch-BPR | 0.700 | 0.418 | 0.877 | 0.551 10 | 11 | 12 | ## The requirements are as follows: 13 | * python==3.6 14 | * pandas==0.24.2 15 | * numpy==1.16.2 16 | * pytorch==1.0.1 17 | * tensorboardX==1.6 (mainly useful when you want to visulize the loss, see https://github.com/lanpa/tensorboard-pytorch) 18 | 19 | ## Example to run: 20 | ``` 21 | python main.py --factor_num=16 --lamda=0.001 22 | ``` 23 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | # dataset name 2 | dataset = 'ml-1m' 3 | assert dataset in ['ml-1m', 'pinterest-20'] 4 | 5 | # paths 6 | main_path = '/home/share/guoyangyang/recommendation/NCF-Data/' 7 | 8 | train_rating = main_path + '{}.train.rating'.format(dataset) 9 | test_rating = main_path + '{}.test.rating'.format(dataset) 10 | test_negative = main_path + '{}.test.negative'.format(dataset) 11 | 12 | model_path = './models/' 13 | BPR_model_path = model_path + 'NeuMF.pth' 14 | -------------------------------------------------------------------------------- /data_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import scipy.sparse as sp 4 | 5 | import torch.utils.data as data 6 | 7 | import config 8 | 9 | 10 | def load_all(test_num=100): 11 | """ We load all the three file here to save time in each epoch. """ 12 | train_data = pd.read_csv( 13 | config.train_rating, 14 | sep='\t', header=None, names=['user', 'item'], 15 | usecols=[0, 1], dtype={0: np.int32, 1: np.int32}) 16 | 17 | user_num = train_data['user'].max() + 1 18 | item_num = train_data['item'].max() + 1 19 | 20 | train_data = train_data.values.tolist() 21 | 22 | # load ratings as a dok matrix 23 | train_mat = sp.dok_matrix((user_num, item_num), dtype=np.float32) 24 | for x in train_data: 25 | train_mat[x[0], x[1]] = 1.0 26 | 27 | test_data = [] 28 | with open(config.test_negative, 'r') as fd: 29 | line = fd.readline() 30 | while line != None and line != '': 31 | arr = line.split('\t') 32 | u = eval(arr[0])[0] 33 | test_data.append([u, eval(arr[0])[1]]) 34 | for i in arr[1:]: 35 | test_data.append([u, int(i)]) 36 | line = fd.readline() 37 | return train_data, test_data, user_num, item_num, train_mat 38 | 39 | 40 | class BPRData(data.Dataset): 41 | def __init__(self, features, 42 | num_item, train_mat=None, num_ng=0, is_training=None): 43 | super(BPRData, self).__init__() 44 | """ Note that the labels are only useful when training, we thus 45 | add them in the ng_sample() function. 46 | """ 47 | self.features = features 48 | self.num_item = num_item 49 | self.train_mat = train_mat 50 | self.num_ng = num_ng 51 | self.is_training = is_training 52 | 53 | def ng_sample(self): 54 | assert self.is_training, 'no need to sampling when testing' 55 | 56 | self.features_fill = [] 57 | for x in self.features: 58 | u, i = x[0], x[1] 59 | for t in range(self.num_ng): 60 | j = np.random.randint(self.num_item) 61 | while (u, j) in self.train_mat: 62 | j = np.random.randint(self.num_item) 63 | self.features_fill.append([u, i, j]) 64 | 65 | def __len__(self): 66 | return self.num_ng * len(self.features) if \ 67 | self.is_training else len(self.features) 68 | 69 | def __getitem__(self, idx): 70 | features = self.features_fill if \ 71 | self.is_training else self.features 72 | 73 | user = features[idx][0] 74 | item_i = features[idx][1] 75 | item_j = features[idx][2] if \ 76 | self.is_training else features[idx][1] 77 | return user, item_i, item_j 78 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | 5 | def hit(gt_item, pred_items): 6 | if gt_item in pred_items: 7 | return 1 8 | return 0 9 | 10 | 11 | def ndcg(gt_item, pred_items): 12 | if gt_item in pred_items: 13 | index = pred_items.index(gt_item) 14 | return np.reciprocal(np.log2(index+2)) 15 | return 0 16 | 17 | 18 | def metrics(model, test_loader, top_k): 19 | HR, NDCG = [], [] 20 | 21 | for user, item_i, item_j in test_loader: 22 | user = user.cuda() 23 | item_i = item_i.cuda() 24 | item_j = item_j.cuda() # not useful when testing 25 | 26 | prediction_i, prediction_j = model(user, item_i, item_j) 27 | _, indices = torch.topk(prediction_i, top_k) 28 | recommends = torch.take( 29 | item_i, indices).cpu().numpy().tolist() 30 | 31 | gt_item = item_i[0].item() 32 | HR.append(hit(gt_item, recommends)) 33 | NDCG.append(ndcg(gt_item, recommends)) 34 | 35 | return np.mean(HR), np.mean(NDCG) 36 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import argparse 4 | import numpy as np 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | import torch.utils.data as data 10 | import torch.backends.cudnn as cudnn 11 | from tensorboardX import SummaryWriter 12 | 13 | import model 14 | import config 15 | import evaluate 16 | import data_utils 17 | 18 | 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument("--lr", 21 | type=float, 22 | default=0.01, 23 | help="learning rate") 24 | parser.add_argument("--lamda", 25 | type=float, 26 | default=0.001, 27 | help="model regularization rate") 28 | parser.add_argument("--batch_size", 29 | type=int, 30 | default=4096, 31 | help="batch size for training") 32 | parser.add_argument("--epochs", 33 | type=int, 34 | default=50, 35 | help="training epoches") 36 | parser.add_argument("--top_k", 37 | type=int, 38 | default=10, 39 | help="compute metrics@top_k") 40 | parser.add_argument("--factor_num", 41 | type=int, 42 | default=32, 43 | help="predictive factors numbers in the model") 44 | parser.add_argument("--num_ng", 45 | type=int, 46 | default=4, 47 | help="sample negative items for training") 48 | parser.add_argument("--test_num_ng", 49 | type=int, 50 | default=99, 51 | help="sample part of negative items for testing") 52 | parser.add_argument("--out", 53 | default=True, 54 | help="save model or not") 55 | parser.add_argument("--gpu", 56 | type=str, 57 | default="0", 58 | help="gpu card ID") 59 | args = parser.parse_args() 60 | 61 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 62 | cudnn.benchmark = True 63 | 64 | 65 | ############################## PREPARE DATASET ########################## 66 | train_data, test_data, user_num ,item_num, train_mat = data_utils.load_all() 67 | 68 | # construct the train and test datasets 69 | train_dataset = data_utils.BPRData( 70 | train_data, item_num, train_mat, args.num_ng, True) 71 | test_dataset = data_utils.BPRData( 72 | test_data, item_num, train_mat, 0, False) 73 | train_loader = data.DataLoader(train_dataset, 74 | batch_size=args.batch_size, shuffle=True, num_workers=4) 75 | test_loader = data.DataLoader(test_dataset, 76 | batch_size=args.test_num_ng+1, shuffle=False, num_workers=0) 77 | 78 | ########################### CREATE MODEL ################################# 79 | model = model.BPR(user_num, item_num, args.factor_num) 80 | model.cuda() 81 | 82 | optimizer = optim.SGD( 83 | model.parameters(), lr=args.lr, weight_decay=args.lamda) 84 | # writer = SummaryWriter() # for visualization 85 | 86 | ########################### TRAINING ##################################### 87 | count, best_hr = 0, 0 88 | for epoch in range(args.epochs): 89 | model.train() 90 | start_time = time.time() 91 | train_loader.dataset.ng_sample() 92 | 93 | for user, item_i, item_j in train_loader: 94 | user = user.cuda() 95 | item_i = item_i.cuda() 96 | item_j = item_j.cuda() 97 | 98 | model.zero_grad() 99 | prediction_i, prediction_j = model(user, item_i, item_j) 100 | loss = - (prediction_i - prediction_j).sigmoid().log().sum() 101 | loss.backward() 102 | optimizer.step() 103 | # writer.add_scalar('data/loss', loss.item(), count) 104 | count += 1 105 | 106 | model.eval() 107 | HR, NDCG = evaluate.metrics(model, test_loader, args.top_k) 108 | 109 | elapsed_time = time.time() - start_time 110 | print("The time elapse of epoch {:03d}".format(epoch) + " is: " + 111 | time.strftime("%H: %M: %S", time.gmtime(elapsed_time))) 112 | print("HR: {:.3f}\tNDCG: {:.3f}".format(np.mean(HR), np.mean(NDCG))) 113 | 114 | if HR > best_hr: 115 | best_hr, best_ndcg, best_epoch = HR, NDCG, epoch 116 | if args.out: 117 | if not os.path.exists(config.model_path): 118 | os.mkdir(config.model_path) 119 | torch.save(model, '{}BPR.pt'.format(config.model_path)) 120 | 121 | print("End. Best epoch {:03d}: HR = {:.3f}, \ 122 | NDCG = {:.3f}".format(best_epoch, best_hr, best_ndcg)) 123 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class BPR(nn.Module): 6 | def __init__(self, user_num, item_num, factor_num): 7 | super(BPR, self).__init__() 8 | """ 9 | user_num: number of users; 10 | item_num: number of items; 11 | factor_num: number of predictive factors. 12 | """ 13 | self.embed_user = nn.Embedding(user_num, factor_num) 14 | self.embed_item = nn.Embedding(item_num, factor_num) 15 | 16 | nn.init.normal_(self.embed_user.weight, std=0.01) 17 | nn.init.normal_(self.embed_item.weight, std=0.01) 18 | 19 | def forward(self, user, item_i, item_j): 20 | user = self.embed_user(user) 21 | item_i = self.embed_item(item_i) 22 | item_j = self.embed_item(item_j) 23 | 24 | prediction_i = (user * item_i).sum(dim=-1) 25 | prediction_j = (user * item_j).sum(dim=-1) 26 | return prediction_i, prediction_j 27 | --------------------------------------------------------------------------------