extend Melu code to perform different meta algorithms and hyperparameters
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.

learnToLearn.py 6.1KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. import os
  2. import torch
  3. import pickle
  4. from MeLU import MeLU
  5. from options import config
  6. from model_training import training
  7. from data_generation import generate
  8. from evidence_candidate import selection
  9. from model_test import test
  10. from embedding_module import EmbeddingModule
  11. import learn2learn as l2l
  12. from embeddings import item, user
  13. import random
  14. import numpy as np
  15. from learnToLearnTest import test
  16. from fast_adapt import fast_adapt
  17. # DATA GENERATION
  18. print("DATA GENERATION PHASE")
  19. os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  20. os.environ["CUDA_VISIBLE_DEVICES"] = "1"
  21. master_path= "/media/external_10TB/10TB/maheri/melu_data5"
  22. if not os.path.exists("{}/".format(master_path)):
  23. os.mkdir("{}/".format(master_path))
  24. # preparing dataset. It needs about 22GB of your hard disk space.
  25. generate(master_path)
  26. # TRAINING
  27. print("TRAINING PHASE")
  28. embedding_dim = config['embedding_dim']
  29. fc1_in_dim = config['embedding_dim'] * 8
  30. fc2_in_dim = config['first_fc_hidden_dim']
  31. fc2_out_dim = config['second_fc_hidden_dim']
  32. use_cuda = config['use_cuda']
  33. emb = EmbeddingModule(config).cuda()
  34. fc1 = torch.nn.Linear(fc1_in_dim, fc2_in_dim)
  35. fc2 = torch.nn.Linear(fc2_in_dim, fc2_out_dim)
  36. linear_out = torch.nn.Linear(fc2_out_dim, 1)
  37. head = torch.nn.Sequential(fc1,fc2,linear_out)
  38. # META LEARNING
  39. print("META LEARNING PHASE")
  40. head = l2l.algorithms.MetaSGD(head, lr=config['local_lr'])
  41. head.cuda()
  42. # Setup optimization
  43. print("SETUP OPTIMIZATION PHASE")
  44. all_parameters = list(emb.parameters()) + list(head.parameters())
  45. optimizer = torch.optim.Adam(all_parameters, lr=config['lr'])
  46. # loss = torch.nn.MSELoss(reduction='mean')
  47. # Load training dataset.
  48. print("LOAD DATASET PHASE")
  49. training_set_size = int(len(os.listdir("{}/warm_state".format(master_path))) / 4)
  50. supp_xs_s = []
  51. supp_ys_s = []
  52. query_xs_s = []
  53. query_ys_s = []
  54. for idx in range(training_set_size):
  55. supp_xs_s.append(pickle.load(open("{}/warm_state/supp_x_{}.pkl".format(master_path, idx), "rb")))
  56. supp_ys_s.append(pickle.load(open("{}/warm_state/supp_y_{}.pkl".format(master_path, idx), "rb")))
  57. query_xs_s.append(pickle.load(open("{}/warm_state/query_x_{}.pkl".format(master_path, idx), "rb")))
  58. query_ys_s.append(pickle.load(open("{}/warm_state/query_y_{}.pkl".format(master_path, idx), "rb")))
  59. total_dataset = list(zip(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s))
  60. del(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s)
  61. training_set_size = len(total_dataset)
  62. batch_size = config['batch_size']
  63. print("\n\n\n")
  64. for iteration in range(config['num_epoch']):
  65. random.shuffle(total_dataset)
  66. num_batch = int(training_set_size / batch_size)
  67. a, b, c, d = zip(*total_dataset)
  68. for i in range(num_batch):
  69. optimizer.zero_grad()
  70. meta_train_error = 0.0
  71. meta_train_accuracy = 0.0
  72. meta_valid_error = 0.0
  73. meta_valid_accuracy = 0.0
  74. meta_test_error = 0.0
  75. meta_test_accuracy = 0.0
  76. print("EPOCH: ", iteration, " BATCH: ", i)
  77. supp_xs = list(a[batch_size * i:batch_size * (i + 1)])
  78. supp_ys = list(b[batch_size * i:batch_size * (i + 1)])
  79. query_xs = list(c[batch_size * i:batch_size * (i + 1)])
  80. query_ys = list(d[batch_size * i:batch_size * (i + 1)])
  81. batch_sz = len(supp_xs)
  82. for j in range(batch_size):
  83. supp_xs[j] = supp_xs[j].cuda()
  84. supp_ys[j] = supp_ys[j].cuda()
  85. query_xs[j] = query_xs[j].cuda()
  86. query_ys[j] = query_ys[j].cuda()
  87. for task in range(batch_sz):
  88. # print("EPOCH: ", iteration," BATCH: ",i, "TASK: ",task)
  89. # Compute meta-training loss
  90. learner = head.clone()
  91. temp_sxs = emb(supp_xs[task])
  92. temp_qxs = emb(query_xs[task])
  93. evaluation_error = fast_adapt(learner,
  94. temp_sxs,
  95. temp_qxs,
  96. supp_ys[task],
  97. query_ys[task],
  98. config['inner']
  99. )
  100. evaluation_error.backward()
  101. meta_train_error += evaluation_error.item()
  102. # Print some metrics
  103. print('Iteration', iteration)
  104. print('Meta Train Error', meta_train_error / batch_sz)
  105. # print('Meta Train Accuracy', meta_train_accuracy / batch_sz)
  106. # print('Meta Valid Error', meta_valid_error / batch_sz)
  107. # print('Meta Valid Accuracy', meta_valid_accuracy / batch_sz)
  108. # print('Meta Test Error', meta_test_error / batch_sz)
  109. # print('Meta Test Accuracy', meta_test_accuracy / batch_sz)
  110. # Average the accumulated gradients and optimize
  111. for p in all_parameters:
  112. p.grad.data.mul_(1.0 / batch_sz)
  113. optimizer.step()
  114. print("===============================================\n")
  115. # save model
  116. final_model = torch.nn.Sequential(emb,head)
  117. torch.save(final_model.state_dict(), master_path + "/models_sgd.pkl")
  118. # testing
  119. print("start of test phase")
  120. for test_state in ['warm_state', 'user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
  121. test_dataset = None
  122. test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
  123. supp_xs_s = []
  124. supp_ys_s = []
  125. query_xs_s = []
  126. query_ys_s = []
  127. for idx in range(test_set_size):
  128. supp_xs_s.append(pickle.load(open("{}/{}/supp_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  129. supp_ys_s.append(pickle.load(open("{}/{}/supp_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  130. query_xs_s.append(pickle.load(open("{}/{}/query_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  131. query_ys_s.append(pickle.load(open("{}/{}/query_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  132. test_dataset = list(zip(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s))
  133. del (supp_xs_s, supp_ys_s, query_xs_s, query_ys_s)
  134. print("===================== " + test_state + " =====================")
  135. test(emb,head, test_dataset, batch_size=config['batch_size'], num_epoch=config['num_epoch'])
  136. print("===================================================\n\n\n")