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

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386
  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. from clustering import ClustringModule, Trainer
  15. import numpy as np
  16. from torch.nn import functional as F
  17. def data_batching(indexes, C_distribs, batch_size, training_set_size, num_clusters):
  18. probs = np.squeeze(C_distribs)
  19. cs = [np.random.choice(num_clusters, p=i) for i in probs]
  20. num_batch = int(training_set_size / batch_size)
  21. res = [[] for i in range(num_batch)]
  22. clas = [[] for i in range(num_clusters)]
  23. for idx, c in zip(indexes, cs):
  24. clas[c].append(idx)
  25. t = np.array([len(i) for i in clas])
  26. t = t / t.sum()
  27. dif = list(set(list(np.arange(training_set_size))) - set(indexes[0:(num_batch * batch_size)]))
  28. cnt = 0
  29. for i in range(len(res)):
  30. for j in range(batch_size):
  31. temp = np.random.choice(num_clusters, p=t)
  32. if len(clas[temp]) > 0:
  33. res[i].append(clas[temp].pop(0))
  34. else:
  35. # res[i].append(indexes[training_set_size-1-cnt])
  36. res[i].append(random.choice(dif))
  37. cnt = cnt + 1
  38. res = np.random.permutation(res)
  39. final_result = np.array(res).flatten()
  40. return final_result
  41. def parse_args():
  42. print("==============")
  43. parser = argparse.ArgumentParser([], description='Fast Context Adaptation via Meta-Learning (CAVIA),'
  44. 'Clasification experiments.')
  45. print("==============\n")
  46. parser.add_argument('--seed', type=int, default=53)
  47. parser.add_argument('--task', type=str, default='multi', help='problem setting: sine or celeba')
  48. parser.add_argument('--tasks_per_metaupdate', type=int, default=32,
  49. help='number of tasks in each batch per meta-update')
  50. parser.add_argument('--lr_inner', type=float, default=5e-6, help='inner-loop learning rate (per task)')
  51. parser.add_argument('--lr_meta', type=float, default=5e-5,
  52. help='outer-loop learning rate (used with Adam optimiser)')
  53. # parser.add_argument('--lr_meta_decay', type=float, default=0.9, help='decay factor for meta learning rate')
  54. parser.add_argument('--inner', type=int, default=1,
  55. help='number of gradient steps in inner loop (during training)')
  56. parser.add_argument('--inner_eval', type=int, default=1,
  57. help='number of gradient updates at test time (for evaluation)')
  58. parser.add_argument('--first_order', action='store_true', default=False,
  59. help='run first order approximation of CAVIA')
  60. parser.add_argument('--adapt_transform', action='store_true', default=False,
  61. help='run adaptation transform')
  62. parser.add_argument('--transformer', type=str, default="kronoker",
  63. help='transformer type')
  64. parser.add_argument('--meta_algo', type=str, default="metasgd",
  65. help='MAML/MetaSGD/GBML')
  66. parser.add_argument('--gpu', type=int, default=0,
  67. help='number of gpu to run the code')
  68. parser.add_argument('--epochs', type=int, default=config['num_epoch'],
  69. help='number of gpu to run the code')
  70. # parser.add_argument('--data_root', type=str, default="./movielens/ml-1m", help='path to data root')
  71. # parser.add_argument('--num_workers', type=int, default=4, help='num of workers to use')
  72. # parser.add_argument('--test', action='store_true', default=False, help='num of workers to use')
  73. # parser.add_argument('--embedding_dim', type=int, default=32, help='num of workers to use')
  74. # parser.add_argument('--first_fc_hidden_dim', type=int, default=64, help='num of workers to use')
  75. # parser.add_argument('--second_fc_hidden_dim', type=int, default=64, help='num of workers to use')
  76. # parser.add_argument('--num_epoch', type=int, default=30, help='num of workers to use')
  77. # parser.add_argument('--num_genre', type=int, default=25, help='num of workers to use')
  78. # parser.add_argument('--num_director', type=int, default=2186, help='num of workers to use')
  79. # parser.add_argument('--num_actor', type=int, default=8030, help='num of workers to use')
  80. # parser.add_argument('--num_rate', type=int, default=6, help='num of workers to use')
  81. # parser.add_argument('--num_gender', type=int, default=2, help='num of workers to use')
  82. # parser.add_argument('--num_age', type=int, default=7, help='num of workers to use')
  83. # parser.add_argument('--num_occupation', type=int, default=21, help='num of workers to use')
  84. # parser.add_argument('--num_zipcode', type=int, default=3402, help='num of workers to use')
  85. # parser.add_argument('--rerun', action='store_true', default=False,
  86. # help='Re-run experiment (will override previously saved results)')
  87. args = parser.parse_args()
  88. # use the GPU if available
  89. # args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  90. # print('Running on device: {}'.format(args.device))
  91. return args
  92. from torch.nn import functional as F
  93. def kl_loss(C_distribs):
  94. # batchsize * k
  95. C_distribs = torch.stack(C_distribs).squeeze()
  96. # print("injam:",len(C_distribs))
  97. # print(C_distribs[0].shape)
  98. # batchsize * k
  99. # print("injam2",C_distribs)
  100. C_distribs_sq = torch.pow(C_distribs, 2)
  101. # print("injam3",C_distribs_sq)
  102. # 1*k
  103. C_distribs_sum = torch.sum(C_distribs, dim=0, keepdim=True)
  104. # print("injam4",C_distribs_sum)
  105. # batchsize * k
  106. temp = C_distribs_sq / C_distribs_sum
  107. # print("injam5",temp)
  108. # batchsize * 1
  109. temp_sum = torch.sum(temp, dim=1, keepdim=True)
  110. # print("injam6",temp_sum)
  111. target_distribs = temp / temp_sum
  112. # print("injam7",target_distribs)
  113. # calculate the kl loss
  114. clustering_loss = F.kl_div(C_distribs.log(), target_distribs, reduction='batchmean')
  115. # print("injam8",clustering_loss)
  116. return clustering_loss
  117. if __name__ == '__main__':
  118. args = parse_args()
  119. print(args)
  120. if config['use_cuda']:
  121. os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  122. os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
  123. master_path = "/media/external_10TB/10TB/maheri/define_task_melu_data2"
  124. config['master_path'] = master_path
  125. # DATA GENERATION
  126. print("DATA GENERATION PHASE")
  127. if not os.path.exists("{}/".format(master_path)):
  128. os.mkdir("{}/".format(master_path))
  129. # preparing dataset. It needs about 22GB of your hard disk space.
  130. generate(master_path)
  131. # TRAINING
  132. print("TRAINING PHASE")
  133. embedding_dim = config['embedding_dim']
  134. fc1_in_dim = config['embedding_dim'] * 8
  135. fc2_in_dim = config['first_fc_hidden_dim']
  136. fc2_out_dim = config['second_fc_hidden_dim']
  137. use_cuda = config['use_cuda']
  138. if use_cuda:
  139. emb = EmbeddingModule(config).cuda()
  140. else:
  141. emb = EmbeddingModule(config)
  142. # META LEARNING
  143. print("META LEARNING PHASE")
  144. # define transformer
  145. transform = None
  146. if args.transformer == "kronoker":
  147. transform = KroneckerTransform(l2l.nn.KroneckerLinear)
  148. elif args.transformer == "linear":
  149. transform = l2l.optim.ModuleTransform(torch.nn.Linear)
  150. trainer = Trainer(config)
  151. tr = trainer
  152. # define meta algorithm
  153. if args.meta_algo == "maml":
  154. trainer = l2l.algorithms.MAML(trainer, lr=args.lr_inner, first_order=args.first_order)
  155. elif args.meta_algo == 'metasgd':
  156. trainer = l2l.algorithms.MetaSGD(trainer, lr=args.lr_inner, first_order=args.first_order)
  157. elif args.meta_algo == 'gbml':
  158. trainer = l2l.algorithms.GBML(trainer, transform=transform, lr=args.lr_inner,
  159. adapt_transform=args.adapt_transform,
  160. first_order=args.first_order)
  161. if use_cuda:
  162. trainer.cuda()
  163. # Setup optimization
  164. print("SETUP OPTIMIZATION PHASE")
  165. all_parameters = list(emb.parameters()) + list(trainer.parameters())
  166. optimizer = torch.optim.Adam(all_parameters, lr=config['lr'])
  167. # loss = torch.nn.MSELoss(reduction='mean')
  168. # Load training dataset.
  169. print("LOAD DATASET PHASE")
  170. training_set_size = int(len(os.listdir("{}/warm_state".format(master_path))) / 4)
  171. supp_xs_s = []
  172. supp_ys_s = []
  173. query_xs_s = []
  174. query_ys_s = []
  175. batch_size = config['batch_size']
  176. # torch.cuda.empty_cache()
  177. print("\n\n\n")
  178. for iteration in range(config['num_epoch']):
  179. if iteration == 0:
  180. print("changing cluster centroids started ...")
  181. indexes = list(np.arange(training_set_size))
  182. supp_xs, supp_ys, query_xs, query_ys = [], [], [], []
  183. for idx in range(0, 2500):
  184. supp_xs.append(pickle.load(open("{}/warm_state/supp_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
  185. supp_ys.append(pickle.load(open("{}/warm_state/supp_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
  186. query_xs.append(
  187. pickle.load(open("{}/warm_state/query_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
  188. query_ys.append(
  189. pickle.load(open("{}/warm_state/query_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
  190. batch_sz = len(supp_xs)
  191. user_embeddings = []
  192. for task in range(batch_sz):
  193. # Compute meta-training loss
  194. supp_xs[task] = supp_xs[task].cuda()
  195. supp_ys[task] = supp_ys[task].cuda()
  196. # query_xs[task] = query_xs[task].cuda()
  197. # query_ys[task] = query_ys[task].cuda()
  198. temp_sxs = emb(supp_xs[task])
  199. # temp_qxs = emb(query_xs[task])
  200. y = supp_ys[task].view(-1, 1)
  201. # input_pairs = torch.cat((temp_sxs, y), dim=1)
  202. input_pairs = temp_sxs
  203. task_embed = tr.cluster_module.input_to_hidden(input_pairs)
  204. # todo : may be useless
  205. mean_task = tr.cluster_module.aggregate(task_embed)
  206. user_embeddings.append(mean_task.detach().cpu().numpy())
  207. supp_xs[task] = supp_xs[task].cpu()
  208. supp_ys[task] = supp_ys[task].cpu()
  209. from sklearn.cluster import KMeans
  210. user_embeddings = np.array(user_embeddings)
  211. kmeans_model = KMeans(n_clusters=config['cluster_k'], init="k-means++").fit(user_embeddings)
  212. tr.cluster_module.array.data = torch.Tensor(kmeans_model.cluster_centers_).cuda()
  213. if iteration > 0:
  214. # indexes = data_batching(indexes, C_distribs, batch_size, training_set_size, config['cluster_k'])
  215. # random.shuffle(indexes)
  216. num_batch = int(training_set_size / batch_size)
  217. indexes = list(np.arange(training_set_size))
  218. random.shuffle(indexes)
  219. else:
  220. num_batch = int(training_set_size / batch_size)
  221. indexes = list(np.arange(training_set_size))
  222. random.shuffle(indexes)
  223. for i in range(num_batch):
  224. meta_train_error = 0.0
  225. meta_cluster_error = 0.0
  226. optimizer.zero_grad()
  227. print("EPOCH: ", iteration, " BATCH: ", i)
  228. supp_xs, supp_ys, query_xs, query_ys = [], [], [], []
  229. for idx in range(batch_size * i, batch_size * (i + 1)):
  230. supp_xs.append(pickle.load(open("{}/warm_state/supp_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
  231. supp_ys.append(pickle.load(open("{}/warm_state/supp_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
  232. query_xs.append(
  233. pickle.load(open("{}/warm_state/query_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
  234. query_ys.append(
  235. pickle.load(open("{}/warm_state/query_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
  236. batch_sz = len(supp_xs)
  237. if use_cuda:
  238. for j in range(batch_size):
  239. supp_xs[j] = supp_xs[j].cuda()
  240. supp_ys[j] = supp_ys[j].cuda()
  241. query_xs[j] = query_xs[j].cuda()
  242. query_ys[j] = query_ys[j].cuda()
  243. C_distribs = []
  244. for task in range(batch_sz):
  245. learner = trainer.clone()
  246. temp_sxs = emb(supp_xs[task])
  247. temp_qxs = emb(query_xs[task])
  248. evaluation_error, c, k_loss = fast_adapt(learner,
  249. temp_sxs,
  250. temp_qxs,
  251. supp_ys[task],
  252. query_ys[task],
  253. config['inner'],
  254. epoch=iteration)
  255. # C_distribs.append(c)
  256. evaluation_error.backward(retain_graph=True)
  257. meta_train_error += evaluation_error.item()
  258. meta_cluster_error += k_loss
  259. # Print some metrics
  260. print('Iteration', iteration)
  261. print('Meta Train Error', meta_train_error / batch_sz)
  262. print('KL Train Error', meta_cluster_error / batch_sz)
  263. # clustering_loss = config['kl_loss_weight'] * kl_loss(C_distribs)
  264. # clustering_loss.backward()
  265. # print("kl_loss:", round(clustering_loss.item(), 8), "\t", C_distribs[0].cpu().detach().numpy())
  266. # Average the accumulated gradients and optimize
  267. for p in all_parameters:
  268. p.grad.data.mul_(1.0 / batch_sz)
  269. optimizer.step()
  270. # torch.cuda.empty_cache()
  271. # del (supp_xs, supp_ys, query_xs, query_ys, learner, temp_sxs, temp_qxs)
  272. # gc.collect()
  273. print("===============================================\n")
  274. if iteration % 2 == 0 and iteration != 0:
  275. # testing
  276. print("start of test phase")
  277. trainer.eval()
  278. with open("results2.txt", "a") as f:
  279. f.write("epoch:{}\n".format(iteration))
  280. for test_state in ['user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
  281. test_dataset = None
  282. test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
  283. supp_xs_s = []
  284. supp_ys_s = []
  285. query_xs_s = []
  286. query_ys_s = []
  287. gc.collect()
  288. print("===================== " + test_state + " =====================")
  289. mse, ndc1, ndc3 = test(emb, trainer, test_dataset, batch_size=config['batch_size'],
  290. num_epoch=config['num_epoch'], test_state=test_state, args=args)
  291. with open("results2.txt", "a") as f:
  292. f.write("{}\t{}\t{}\n".format(mse, ndc1, ndc3))
  293. print("===================================================")
  294. del (test_dataset)
  295. gc.collect()
  296. trainer.train()
  297. with open("results2.txt", "a") as f:
  298. f.write("\n")
  299. print("\n\n\n")
  300. # save model
  301. # final_model = torch.nn.Sequential(emb, head)
  302. # torch.save(final_model.state_dict(), master_path + "/models_gbml.pkl")
  303. # testing
  304. # print("start of test phase")
  305. # for test_state in ['warm_state', 'user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
  306. # test_dataset = None
  307. # test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
  308. # supp_xs_s = []
  309. # supp_ys_s = []
  310. # query_xs_s = []
  311. # query_ys_s = []
  312. # for idx in range(test_set_size):
  313. # supp_xs_s.append(pickle.load(open("{}/{}/supp_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  314. # supp_ys_s.append(pickle.load(open("{}/{}/supp_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  315. # query_xs_s.append(pickle.load(open("{}/{}/query_x_{}.pkl".format(master_path, test_state, idx), "rb")))
  316. # query_ys_s.append(pickle.load(open("{}/{}/query_y_{}.pkl".format(master_path, test_state, idx), "rb")))
  317. # test_dataset = list(zip(supp_xs_s, supp_ys_s, query_xs_s, query_ys_s))
  318. # del (supp_xs_s, supp_ys_s, query_xs_s, query_ys_s)
  319. #
  320. # print("===================== " + test_state + " =====================")
  321. # test(emb, head, test_dataset, batch_size=config['batch_size'], num_epoch=args.epochs,
  322. # adaptation_step=args.inner_eval)
  323. # print("===================================================\n\n\n")
  324. # print(args)