def main(Epoch=5, Pretrain=False, Predict=False, Batch_size=2, log=None):
print('*'*150)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)
print("Get Dataset")
dataset_4qs = pickle.load(open('./data/tweets-cikm.txtfoursquare.pk', 'rb'))
print("User number: ", len(dataset_4qs['uid_list']))
print("Generate Traj_list")
if Predict:
train_input_list, test_input_list = generate_input2(dataset_4qs['data_neural'])
else:
traj_input_list = generate_input(dataset_4qs['data_neural'])
print("Generate Train Dataset")
if Predict:
train_dataset = Predict_Traj_Dataset(data_neural=dataset_4qs['data_neural'], traj_list=train_input_list, vid_list=dataset_4qs['vid_list'])
valid_dataset = Predict_Traj_Dataset(data_neural=dataset_4qs['data_neural'], traj_list=test_input_list, vid_list=dataset_4qs['vid_list'])
print("Generate Train Dataloader")
train_traj_dataloader = DataLoader(train_dataset, batch_size=Batch_size, shuffle=True)
valid_traj_dataloader = DataLoader(valid_dataset, batch_size=Batch_size, shuffle=True)
else:
train_dataset = Bert_Traj_Dataset(data_neural=dataset_4qs['data_neural'], traj_list=traj_input_list, vid_list=dataset_4qs['vid_list'])
print("Generate Train Dataloader")
train_traj_dataloader = DataLoader(train_dataset, batch_size=Batch_size,shuffle=True)
if Predict:
if Pretrain:
print("Loaded Pretrained Bert")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1)
Bert.load_state_dict(torch.load('./pretrain/bert_trained_ep14.pth'))
else:
print("Create New Bert")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1)
print("Get Predict Model")
model = Predict_Model(Bert, token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.2)
# model.load_state_dict(torch.load('./pretrain/Predict_model_trained_ep45.pth'))
model = model.to(device)
else:
if Pretrain:
print("Get Pretrained Model")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1)
model = BertLM(token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1)
model.load_state_dict(torch.load('./pretrain/bert_trained'))
model = model.to(device)
else:
print("Initialite Model")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1)
model = BertLM(Bert_Traj_Model=Bert, token_size=len(dataset_4qs['vid_list']), head_n=12, d_model=512, N_layers=12, dropout=0.1).to(device)
# optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
optimizer = Trans_Optim(
torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),
init_lr=2, d_model=512, n_warmup_steps=10000)
print('*'*150)
print('-'*65 + " START TRAIN " + '-'*65)
run(Epoch, model, optimizer, device, train_traj_dataloader, valid_traj_dataloader=valid_traj_dataloader,Predict=Predict, log=log)
def main(Epoch=200,
Bert_Pretrain=False,
Batch_size=10,
Pretrained=False,
log='predict'):
head_n = 5
d_model = 500
N_layers = 10
dropout = 0.3
# head_n = 10
# d_model = 400
# N_layers = 10
# dropout = 0.1
print('*' * 150)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)
print("Get Dataset")
dataset_4qs = pickle.load(open('./data/tweets-cikm.txtfoursquare.pk',
'rb'))
print("User number: ", len(dataset_4qs['uid_list']))
print("Generate Train_traj_list")
train_data, train_traj_idxx = generate_input_history(
data_neural=dataset_4qs["data_neural"], mode="train")
print("Generate Valid_traj_list")
test_data, test_traj_idxx = generate_input_history(
data_neural=dataset_4qs["data_neural"], mode="test")
if Bert_Pretrain:
print("Loaded Pretrained Bert")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']),
head_n=head_n,
d_model=d_model,
N_layers=N_layers,
dropout=dropout)
Bert.load_state_dict(torch.load('./pretrain/bert_trained_ep14.pth'))
else:
print("Create New Bert")
Bert = Bert_Traj_Model(token_size=len(dataset_4qs['vid_list']),
head_n=head_n,
d_model=d_model,
N_layers=N_layers,
dropout=dropout)
print("Get Predict Model")
model = Predict_Model(Bert,
token_size=len(dataset_4qs['vid_list']),
head_n=head_n,
d_model=d_model,
N_layers=N_layers,
dropout=dropout)
if Pretrained:
print("Load Pretrained Predict Model")
model.load_state_dict(
torch.load('./pretrain/Predict_model_trained_ep34.pth'))
model = model.to(device)
# optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
optimizer = Trans_Optim(torch.optim.Adam(model.parameters(),
betas=(0.9, 0.98),
eps=1e-09),
init_lr=2,
d_model=768,
n_warmup_steps=4000)
print('*' * 150)
print('-' * 65 + " START TRAIN " + '-' * 65)
run(Epoch, model, optimizer, device, train_data, train_traj_idxx,
test_data, test_traj_idxx, log, Batch_size)