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.

hyper_tunning.py 2.6KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import os
  2. import torch
  3. import torch.nn as nn
  4. from ray import tune
  5. import pickle
  6. import random
  7. import gc
  8. from trainer import Trainer
  9. import numpy as np
  10. from utils import *
  11. from sampler import *
  12. import os
  13. def train_metatl(conf,checkpoint_dir=None):
  14. SEED = conf["params"]['seed']
  15. torch.manual_seed(SEED)
  16. torch.cuda.manual_seed(SEED)
  17. torch.backends.cudnn.deterministic = True
  18. np.random.seed(SEED)
  19. random.seed(SEED)
  20. params = conf['params']
  21. user_train, usernum_train, itemnum, user_input_test, user_test, user_input_valid, user_valid = data_load(params['dataset'], params['K'])
  22. sampler = WarpSampler(user_train, usernum_train, itemnum, batch_size=params['batch_size'], maxlen=params['K'], n_workers=1)
  23. sampler_test = DataLoader(user_input_test, user_test, itemnum, params)
  24. sampler_valid = DataLoader(user_input_valid, user_valid, itemnum, params)
  25. ps = {
  26. "batch_size" : conf["params"]['batch_size'],
  27. "learning_rate" : conf['learning_rate'],
  28. "epoch" : conf["params"]['epoch'],
  29. "beta" : conf['beta'],
  30. "embed_dim" : conf['embed_dim'],
  31. "margin" : conf['margin'],
  32. "K" : conf["params"]['K'],
  33. }
  34. trainer = Trainer([sampler, sampler_valid, sampler_test], conf["itemnum"], ps)
  35. # trainer.train()
  36. if checkpoint_dir:
  37. print("===================== using checkpoint =====================")
  38. model_state, optimizer_state = torch.load(
  39. os.path.join(checkpoint_dir, "checkpoint"))
  40. trainer.MetaTL.load_state_dict(model_state)
  41. trainer.optimizer.load_state_dict(optimizer_state)
  42. for epoch in range(int(ps['epoch']/1000)):
  43. for e in range(1000):
  44. # sample one batch from data_loader
  45. train_task, curr_rel = trainer.train_data_loader.next_batch()
  46. loss, _, _ = trainer.do_one_step(train_task, iseval=False, curr_rel=curr_rel)
  47. # do evaluation on specific epoch
  48. valid_data = trainer.eval(istest=False, epoch=(-1))
  49. # print('Epoch {} Testing...'.format(e))
  50. # test_data = self.eval(istest=True, epoch=e)
  51. with tune.checkpoint_dir(epoch) as checkpoint_dir:
  52. path = os.path.join(checkpoint_dir, "checkpoint")
  53. torch.save((trainer.MetaTL.state_dict(), trainer.optimizer.state_dict()), path)
  54. tune.report(
  55. MRR=valid_data["MRR"], NDCG10=valid_data['NDCG@10'], NDCG5=valid_data["NDCG@5"], NDCG1=valid_data["NDCG@1"],
  56. Hits10=valid_data["Hits@10"], Hits5=valid_data["Hits@5"], Hits1=valid_data["Hits@1"],
  57. training_iteration=epoch*1000
  58. )