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 12KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266
  1. import os
  2. import torch
  3. import pickle
  4. from options import config
  5. from data_generation import generate
  6. from embedding_module import EmbeddingModule
  7. import learn2learn as l2l
  8. import random
  9. from learnToLearnTest import test
  10. from fast_adapt import fast_adapt
  11. import gc
  12. from learn2learn.optim.transforms import KroneckerTransform
  13. import argparse
  14. def parse_args():
  15. print("==============")
  16. parser = argparse.ArgumentParser([], description='Fast Context Adaptation via Meta-Learning (CAVIA),'
  17. 'Clasification experiments.')
  18. print("==============\n")
  19. parser.add_argument('--seed', type=int, default=53)
  20. parser.add_argument('--task', type=str, default='multi', help='problem setting: sine or celeba')
  21. parser.add_argument('--tasks_per_metaupdate', type=int, default=32,
  22. help='number of tasks in each batch per meta-update')
  23. parser.add_argument('--lr_inner', type=float, default=5e-6, help='inner-loop learning rate (per task)')
  24. parser.add_argument('--lr_meta', type=float, default=5e-5,
  25. help='outer-loop learning rate (used with Adam optimiser)')
  26. # parser.add_argument('--lr_meta_decay', type=float, default=0.9, help='decay factor for meta learning rate')
  27. parser.add_argument('--inner', type=int, default=5,
  28. help='number of gradient steps in inner loop (during training)')
  29. parser.add_argument('--inner_eval', type=int, default=5,
  30. help='number of gradient updates at test time (for evaluation)')
  31. parser.add_argument('--first_order', action='store_true', default=False,
  32. help='run first order approximation of CAVIA')
  33. parser.add_argument('--adapt_transform', action='store_true', default=False,
  34. help='run adaptation transform')
  35. parser.add_argument('--transformer', type=str, default="kronoker",
  36. help='transformer type')
  37. parser.add_argument('--meta_algo', type=str, default="gbml",
  38. help='MAML/MetaSGD/GBML')
  39. parser.add_argument('--gpu', type=int, default=0,
  40. help='number of gpu to run the code')
  41. # parser.add_argument('--data_root', type=str, default="./movielens/ml-1m", help='path to data root')
  42. # parser.add_argument('--num_workers', type=int, default=4, help='num of workers to use')
  43. # parser.add_argument('--test', action='store_true', default=False, help='num of workers to use')
  44. # parser.add_argument('--embedding_dim', type=int, default=32, help='num of workers to use')
  45. # parser.add_argument('--first_fc_hidden_dim', type=int, default=64, help='num of workers to use')
  46. # parser.add_argument('--second_fc_hidden_dim', type=int, default=64, help='num of workers to use')
  47. # parser.add_argument('--num_epoch', type=int, default=30, help='num of workers to use')
  48. # parser.add_argument('--num_genre', type=int, default=25, help='num of workers to use')
  49. # parser.add_argument('--num_director', type=int, default=2186, help='num of workers to use')
  50. # parser.add_argument('--num_actor', type=int, default=8030, help='num of workers to use')
  51. # parser.add_argument('--num_rate', type=int, default=6, help='num of workers to use')
  52. # parser.add_argument('--num_gender', type=int, default=2, help='num of workers to use')
  53. # parser.add_argument('--num_age', type=int, default=7, help='num of workers to use')
  54. # parser.add_argument('--num_occupation', type=int, default=21, help='num of workers to use')
  55. # parser.add_argument('--num_zipcode', type=int, default=3402, help='num of workers to use')
  56. # parser.add_argument('--rerun', action='store_true', default=False,
  57. # help='Re-run experiment (will override previously saved results)')
  58. args = parser.parse_args()
  59. # use the GPU if available
  60. # args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  61. # print('Running on device: {}'.format(args.device))
  62. return args
  63. if __name__ == '__main__':
  64. args = parse_args()
  65. print(args)
  66. if config['use_cuda']:
  67. os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  68. os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
  69. master_path= "/media/external_10TB/10TB/maheri/melu_data5"
  70. # DATA GENERATION
  71. print("DATA GENERATION PHASE")
  72. if not os.path.exists("{}/".format(master_path)):
  73. os.mkdir("{}/".format(master_path))
  74. # preparing dataset. It needs about 22GB of your hard disk space.
  75. generate(master_path)
  76. # TRAINING
  77. print("TRAINING PHASE")
  78. embedding_dim = config['embedding_dim']
  79. fc1_in_dim = config['embedding_dim'] * 8
  80. fc2_in_dim = config['first_fc_hidden_dim']
  81. fc2_out_dim = config['second_fc_hidden_dim']
  82. use_cuda = config['use_cuda']
  83. fc1 = torch.nn.Linear(fc1_in_dim, fc2_in_dim)
  84. fc2 = torch.nn.Linear(fc2_in_dim, fc2_out_dim)
  85. linear_out = torch.nn.Linear(fc2_out_dim, 1)
  86. head = torch.nn.Sequential(fc1,fc2,linear_out)
  87. if use_cuda:
  88. emb = EmbeddingModule(config).cuda()
  89. else:
  90. emb = EmbeddingModule(config)
  91. # META LEARNING
  92. print("META LEARNING PHASE")
  93. # head = l2l.algorithms.MetaSGD(head, lr=config['local_lr'],first_order=True)
  94. # define transformer
  95. transform = None
  96. if args.transformer == "kronoker":
  97. transform = KroneckerTransform(l2l.nn.KroneckerLinear)
  98. elif args.transformer == "linear":
  99. transform = l2l.optim.ModuleTransform(torch.nn.Linear)
  100. # define meta algorithm
  101. if args.meta_algo == "maml":
  102. head = l2l.algorithms.MAML(head, lr=config['local_lr'],first_order=args.first_order)
  103. elif args.meta_algo == 'metasgd':
  104. head = l2l.algorithms.MetaSGD(head, lr=config['local_lr'],first_order=args.first_order)
  105. elif args.meta_algo == 'gbml':
  106. head = l2l.algorithms.GBML(head, transform=transform, lr=config['local_lr'],adapt_transform=args.adapt_transform, first_order=args.first_order)
  107. if use_cuda:
  108. head.cuda()
  109. # Setup optimization
  110. print("SETUP OPTIMIZATION PHASE")
  111. all_parameters = list(emb.parameters()) + list(head.parameters())
  112. optimizer = torch.optim.Adam(all_parameters, lr=config['lr'])
  113. # loss = torch.nn.MSELoss(reduction='mean')
  114. # Load training dataset.
  115. print("LOAD DATASET PHASE")
  116. training_set_size = int(len(os.listdir("{}/warm_state".format(master_path))) / 4)
  117. supp_xs_s = []
  118. supp_ys_s = []
  119. query_xs_s = []
  120. query_ys_s = []
  121. for idx in range(training_set_size):
  122. supp_xs_s.append(pickle.load(open("{}/warm_state/supp_x_{}.pkl".format(master_path, idx), "rb")))
  123. supp_ys_s.append(pickle.load(open("{}/warm_state/supp_y_{}.pkl".format(master_path, idx), "rb")))
  124. query_xs_s.append(pickle.load(open("{}/warm_state/query_x_{}.pkl".format(master_path, idx), "rb")))
  125. query_ys_s.append(pickle.load(open("{}/warm_state/query_y_{}.pkl".format(master_path, idx), "rb")))
  126. total_dataset = list(zip(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s))
  127. del(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s)
  128. training_set_size = len(total_dataset)
  129. batch_size = config['batch_size']
  130. # torch.cuda.empty_cache()
  131. random.shuffle(total_dataset)
  132. num_batch = int(training_set_size / batch_size)
  133. a, b, c, d = zip(*total_dataset)
  134. print("\n\n\n")
  135. for iteration in range(config['num_epoch']):
  136. for i in range(num_batch):
  137. optimizer.zero_grad()
  138. meta_train_error = 0.0
  139. meta_train_accuracy = 0.0
  140. meta_valid_error = 0.0
  141. meta_valid_accuracy = 0.0
  142. meta_test_error = 0.0
  143. meta_test_accuracy = 0.0
  144. print("EPOCH: ", iteration, " BATCH: ", i)
  145. supp_xs = list(a[batch_size * i:batch_size * (i + 1)])
  146. supp_ys = list(b[batch_size * i:batch_size * (i + 1)])
  147. query_xs = list(c[batch_size * i:batch_size * (i + 1)])
  148. query_ys = list(d[batch_size * i:batch_size * (i + 1)])
  149. batch_sz = len(supp_xs)
  150. # if use_cuda:
  151. # for j in range(batch_size):
  152. # supp_xs[j] = supp_xs[j].cuda()
  153. # supp_ys[j] = supp_ys[j].cuda()
  154. # query_xs[j] = query_xs[j].cuda()
  155. # query_ys[j] = query_ys[j].cuda()
  156. for task in range(batch_sz):
  157. # print("EPOCH: ", iteration," BATCH: ",i, "TASK: ",task)
  158. # Compute meta-training loss
  159. # if use_cuda:
  160. sxs = supp_xs[task].cuda()
  161. qxs = query_xs[task].cuda()
  162. sys = supp_ys[task].cuda()
  163. qys = query_ys[task].cuda()
  164. learner = head.clone()
  165. temp_sxs = emb(sxs)
  166. temp_qxs = emb(qxs)
  167. evaluation_error = fast_adapt(learner,
  168. temp_sxs,
  169. temp_qxs,
  170. sys,
  171. qys,
  172. args.inner)
  173. # config['inner'])
  174. evaluation_error.backward()
  175. meta_train_error += evaluation_error.item()
  176. del(sxs,qxs,sys,qys)
  177. supp_xs[task].cpu()
  178. query_xs[task].cpu()
  179. supp_ys[task].cpu()
  180. query_ys[task].cpu()
  181. # Print some metrics
  182. print('Iteration', iteration)
  183. print('Meta Train Error', meta_train_error / batch_sz)
  184. # print('Meta Train Accuracy', meta_train_accuracy / batch_sz)
  185. # print('Meta Valid Error', meta_valid_error / batch_sz)
  186. # print('Meta Valid Accuracy', meta_valid_accuracy / batch_sz)
  187. # print('Meta Test Error', meta_test_error / batch_sz)
  188. # print('Meta Test Accuracy', meta_test_accuracy / batch_sz)
  189. # Average the accumulated gradients and optimize
  190. for p in all_parameters:
  191. p.grad.data.mul_(1.0 / batch_sz)
  192. optimizer.step()
  193. # torch.cuda.empty_cache()
  194. del(supp_xs,supp_ys,query_xs,query_ys)
  195. gc.collect()
  196. print("===============================================\n")
  197. # save model
  198. final_model = torch.nn.Sequential(emb,head)
  199. torch.save(final_model.state_dict(), master_path + "/models_gbml.pkl")
  200. # testing
  201. print("start of test phase")
  202. for test_state in ['warm_state', 'user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
  203. test_dataset = None
  204. test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
  205. supp_xs_s = []
  206. supp_ys_s = []
  207. query_xs_s = []
  208. query_ys_s = []
  209. for idx in range(test_set_size):
  210. supp_xs_s.append(pickle.load(open("{}/{}/supp_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  211. supp_ys_s.append(pickle.load(open("{}/{}/supp_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  212. query_xs_s.append(pickle.load(open("{}/{}/query_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  213. query_ys_s.append(pickle.load(open("{}/{}/query_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  214. test_dataset = list(zip(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s))
  215. del (supp_xs_s, supp_ys_s, query_xs_s, query_ys_s)
  216. print("===================== " + test_state + " =====================")
  217. test(emb,head, test_dataset, batch_size=config['batch_size'], num_epoch=config['num_epoch'],adaptation_step=args.inner_eval)
  218. print("===================================================\n\n\n")
  219. print(args)