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)