Sequential Recommendation for cold-start users with meta transitional learning
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

main.py 2.1KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. from trainer import *
  2. from utils import *
  3. from sampler import *
  4. import json
  5. import argparse
  6. def get_params():
  7. args = argparse.ArgumentParser()
  8. args.add_argument("-data", "--dataset", default="electronics", type=str)
  9. args.add_argument("-seed", "--seed", default=None, type=int)
  10. args.add_argument("-K", "--K", default=3, type=int) #NUMBER OF SHOT
  11. args.add_argument("-dim", "--embed_dim", default=100, type=int)
  12. args.add_argument("-bs", "--batch_size", default=1024, type=int)
  13. args.add_argument("-lr", "--learning_rate", default=0.001, type=float)
  14. args.add_argument("-epo", "--epoch", default=100000, type=int)
  15. args.add_argument("-prt_epo", "--print_epoch", default=100, type=int)
  16. args.add_argument("-eval_epo", "--eval_epoch", default=1000, type=int)
  17. args.add_argument("-b", "--beta", default=5, type=float)
  18. args.add_argument("-m", "--margin", default=1, type=float)
  19. args.add_argument("-p", "--dropout_p", default=0.5, type=float)
  20. args.add_argument("-gpu", "--device", default=1, type=int)
  21. args = args.parse_args()
  22. params = {}
  23. for k, v in vars(args).items():
  24. params[k] = v
  25. params['device'] = torch.device('cuda:'+str(args.device))
  26. # params['device'] = torch.device('cpu')
  27. print("gpu:",params['device'])
  28. return params, args
  29. if __name__ == '__main__':
  30. params, args = get_params()
  31. if params['seed'] is not None:
  32. SEED = params['seed']
  33. torch.manual_seed(SEED)
  34. torch.cuda.manual_seed(SEED)
  35. torch.backends.cudnn.deterministic = True
  36. np.random.seed(SEED)
  37. random.seed(SEED)
  38. user_train, usernum_train, itemnum, user_input_test, user_test, user_input_valid, user_valid = data_load(args.dataset, args.K)
  39. sampler = WarpSampler(user_train, usernum_train, itemnum, batch_size=args.batch_size, maxlen=args.K, n_workers=3)
  40. sampler_test = DataLoader(user_input_test, user_test, itemnum, params)
  41. sampler_valid = DataLoader(user_input_valid, user_valid, itemnum, params)
  42. trainer = Trainer([sampler, sampler_valid, sampler_test], itemnum, params)
  43. trainer.train()
  44. sampler.close()