-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
153 lines (123 loc) · 4.78 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import pandas as pd
import torch
import argparse
from torch.utils.data import DataLoader
import torch.optim as optim
import matplotlib.pyplot as plt
from utils import MovieLens,Download
from model.MLP import MLP
from model.GMF import GMF
from model.NeuMF import NeuMF
from train import Train
from evaluation import metrics
import os
import numpy as np
import time
from parser import args
# check device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'device: {device}')
# print GPU information
if torch.cuda.is_available():
print('Current cuda device:', torch.cuda.current_device())
print('Count of using GPUs:', torch.cuda.device_count())
# print selected model
print(f'model name: {args.model}')
# argparse doesn't supprot boolean type
use_downlaod = True if args.download=='True' else False
use_pretrain = True if args.use_pretrain=='True' else False
save_model = True if args.save_model == 'True' else False
pretrain_dir = 'pretrain'
if not os.path.isdir(pretrain_dir):
os.makedirs(pretrain_dir)
# root path for dataset
root_path='dataset'+args.file_size
print(f'file size:{args.file_size}')
# load dataframe
dataset = Download(root=root_path,file_size=args.file_size,download=use_downlaod)
total_dataframe, train_dataframe, test_dataframe = dataset.split_train_test()
# make torch.utils.data.Data object
train_set = MovieLens(df=train_dataframe,total_df=total_dataframe,ng_ratio=4)
test_set = MovieLens(df=test_dataframe,total_df=total_dataframe,ng_ratio=99)
# get number of unique userID, unique movieID
max_num_users,max_num_items = total_dataframe['userId'].max()+1, total_dataframe['movieId'].max()+1
print('data loaded!')
# dataloader for train_dataset
dataloader_train= DataLoader(dataset=train_set,
batch_size=args.batch,
shuffle=True,
num_workers=0,
)
# dataloader for test_dataset
dataloader_test = DataLoader(dataset=test_set,
batch_size=100,
shuffle=False,
num_workers=0,
drop_last=True
)
# select model among these models ['MLP', 'GMF', 'NeuMF']
if args.model=='MLP':
model = MLP(num_users=max_num_users,
num_items=max_num_items,
num_factor=args.factor,
layer=args.layer,
use_pretrain=use_pretrain,
use_NeuMF=False)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
elif args.model=='GMF':
model = GMF(num_users=args.batch*max_num_users,
num_items=args.batch*max_num_items,
num_factor=args.factor,
use_pretrain=use_pretrain,
use_NeuMF=False)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
elif args.model=='NeuMF':
if use_pretrain:
GMF_dir = os.path.join(pretrain_dir,'GMF.pth')
MLP_dir = os.path.join(pretrain_dir,'MLP.pth')
pretrained_GMF = torch.load(GMF_dir)
pretrained_MLP = torch.load(MLP_dir)
for param in pretrained_GMF.parameters():
param.requires_grad = False
for param in pretrained_MLP.parameters():
param.requires_grad = False
else:
pretrained_GMF = None
pretrained_MLP = None
model = NeuMF(num_users=args.batch*max_num_users,
num_items=args.batch*max_num_items,
num_factor=args.factor,
use_pretrain=use_pretrain,
layer=args.layer,
pretrained_GMF=pretrained_GMF,
pretrained_MLP=pretrained_MLP)
if not use_pretrain :
optimizer = optim.Adam(model.parameters(), lr=args.lr)
else:
optimizer = optim.SGD(model.parameters(),lr=args.lr)
# for parallel GPU
# if torch.cuda.device_count() >1:
# print("Multi gpu", torch.cuda.device_count())
# model = torch.nn.DataParallel(model)
model.to(device)
# objective function is log loss (Cross-entropy loss)
criterion = torch.nn.BCELoss()
if __name__=='__main__' :
train = Train(model=model,
optimizer=optimizer,
criterion=criterion,
epochs=args.epoch,
test_obj=dataloader_test,
dataloader=dataloader_train,
device=device,
print_cost=True,)
# measuring time
start = time.time()
train.train()
if save_model:
pretrain_model_dir = os.path.join(pretrain_dir,args.model+'.pth')
torch.save(model,pretrain_model_dir)
end = time.time()
print(f'training time:{end-start:.5f}')
HR,NDCG = metrics(model,test_loader=dataloader_test,top_k=args.topk,device=device)
print("HR: {:.3f}\tNDCG: {:.3f}".format(np.mean(HR), np.mean(NDCG)))