Revert "deep kmeans loss"

This reverts commit b4d130e2

clustering.py

         # 1*k, k*d, 1*d
         value = torch.mm(C, self.array)
         # simple add operation
-        new_task_embed = value + mean_task
-
-        return C, new_task_embed
+        # new_task_embed = value + mean_task
+        # new_task_embed = value
+        new_task_embed = mean_task
+
+        # print("injam1:", new_task_embed)
+        # print("injam2:", self.array)
+        list_dist = []
+        # list_dist = torch.norm(new_task_embed - self.array, p=2, dim=1,keepdim=True)
+        list_dist = torch.sum(torch.pow(new_task_embed - self.array,2),dim=1)
+        stack_dist = list_dist
+
+        # print("injam3:", stack_dist)
+
+        ## Second, find the minimum squared distance for softmax normalization
+        min_dist = min(list_dist)
+        # print("injam4:", min_dist)
+
+        ## Third, compute exponentials shifted with min_dist to avoid underflow (0/0) issues in softmaxes
+        alpha = config['kmeans_alpha']  # Placeholder tensor for alpha
+        list_exp = []
+        for i in range(self.clusters_k):
+            exp = torch.exp(-alpha * (stack_dist[i] - min_dist))
+            list_exp.append(exp)
+        stack_exp = torch.stack(list_exp)
+        sum_exponentials = torch.sum(stack_exp)
+
+        # print("injam5:", stack_exp, sum_exponentials)
+
+        ## Fourth, compute softmaxes and the embedding/representative distances weighted by softmax
+        list_softmax = []
+        list_weighted_dist = []
+        for j in range(self.clusters_k):
+            softmax = stack_exp[j] / sum_exponentials
+            weighted_dist = stack_dist[j] * softmax
+            list_softmax.append(softmax)
+            list_weighted_dist.append(weighted_dist)
+        stack_weighted_dist = torch.stack(list_weighted_dist)
+
+        kmeans_loss = torch.sum(stack_weighted_dist, dim=0)
+        return C, new_task_embed, kmeans_loss
 
 
 class Trainer(torch.nn.Module):
     def aggregate(self, z_i):
         return torch.mean(z_i, dim=0)
 
-    def forward(self, task_embed, y, training,adaptation_data=None,adaptation_labels=None):
+    def forward(self, task_embed, y, training, adaptation_data=None, adaptation_labels=None):
         if training:
-            C, clustered_task_embed = self.cluster_module(task_embed, y)
+            C, clustered_task_embed, k_loss = self.cluster_module(task_embed, y)
             # hidden layers
             # todo : adding activation function or remove it
             hidden_1 = self.fc1(task_embed)
             y_pred = self.linear_out(hidden_3)
 
         else:
-            C, clustered_task_embed = self.cluster_module(adaptation_data, adaptation_labels)
+            C, clustered_task_embed, k_loss = self.cluster_module(adaptation_data, adaptation_labels)
             beta_1 = torch.tanh(self.film_layer_1_beta(clustered_task_embed))
             gamma_1 = torch.tanh(self.film_layer_1_gamma(clustered_task_embed))
             beta_2 = torch.tanh(self.film_layer_2_beta(clustered_task_embed))
 
             y_pred = self.linear_out(hidden_3)
 
-        return y_pred
+        return y_pred, C, k_loss

fast_adapt.py

-# import torch
-# import pickle
-# from options import config
-# import random
-#
-#
-# def cl_loss(c):
-#     alpha = config['alpha']
-#     beta = config['beta']
-#     d = config['d']
-#     a = torch.div(1,
-#                   torch.add(1, torch.exp(torch.mul(-1, torch.mul(alpha, torch.sub(torch.mul(d, c.squeeze()), beta))))))
-#     # a = 1 / (1 + torch.exp((-1) * alpha * (d * c - beta)))
-#     b = torch.mul(a, torch.mul(torch.sub(1, a), torch.sub(1, torch.mul(2, a))))
-#     # b = 1 * a * (1 - a) * (1 - 2 * a)
-#     loss = torch.sum(b)
-#     return loss
-#
-#
-# def fast_adapt(
-#         learn,
-#         adaptation_data,
-#         evaluation_data,
-#         adaptation_labels,
-#         evaluation_labels,
-#         adaptation_steps,
-#         get_predictions=False,
-#         epoch=None):
-#     is_print = random.random() < 0.05
-#
-#     for step in range(adaptation_steps):
-#         temp, c, k_loss = learn(adaptation_data, adaptation_labels, training=True)
-#         train_error = torch.nn.functional.mse_loss(temp.view(-1), adaptation_labels)
-#         # cluster_loss = cl_loss(c)
-#         # total_loss = train_error + config['cluster_loss_weight'] * cluster_loss
-#         total_loss = train_error + config['kmeans_loss_weight'] * k_loss
-#         learn.adapt(total_loss)
-#         if is_print:
-#             # print("in support:\t", round(k_loss.item(),4))
-#             pass
-#
-#     predictions, c, k_loss = learn(evaluation_data, None, training=False, adaptation_data=adaptation_data,
-#                                    adaptation_labels=adaptation_labels)
-#     valid_error = torch.nn.functional.mse_loss(predictions.view(-1), evaluation_labels)
-#     # cluster_loss = cl_loss(c)
-#     # total_loss = valid_error + config['cluster_loss_weight'] * cluster_loss
-#     total_loss = valid_error + config['kmeans_loss_weight'] * k_loss
-#
-#     if is_print:
-#
-#         # print("in query:\t", round(k_loss.item(),4))
-#         print(c[0].detach().cpu().numpy(),"\t",round(k_loss.item(),3),"\n")
-#
-#     # if random.random() < 0.05:
-#     #     print("cl:", round(cluster_loss.item()), "\t c:", c[0].cpu().data.numpy())
-#
-#     if get_predictions:
-#         return total_loss, predictions
-#     return total_loss, c, k_loss.item()
-
 import torch
 import pickle
+from options import config
+import random
+
+
+def cl_loss(c):
+    alpha = config['alpha']
+    beta = config['beta']
+    d = config['d']
+    a = torch.div(1,
+                  torch.add(1, torch.exp(torch.mul(-1, torch.mul(alpha, torch.sub(torch.mul(d, c.squeeze()), beta))))))
+    # a = 1 / (1 + torch.exp((-1) * alpha * (d * c - beta)))
+    b = torch.mul(a, torch.mul(torch.sub(1, a), torch.sub(1, torch.mul(2, a))))
+    # b = 1 * a * (1 - a) * (1 - 2 * a)
+    loss = torch.sum(b)
+    return loss
 
 
 def fast_adapt(
         adaptation_labels,
         evaluation_labels,
         adaptation_steps,
-        get_predictions=False):
+        get_predictions=False,
+        epoch=None):
+    is_print = random.random() < 0.05
+
     for step in range(adaptation_steps):
-        temp = learn(adaptation_data, adaptation_labels, training=True)
+        temp, c, k_loss = learn(adaptation_data, adaptation_labels, training=True)
         train_error = torch.nn.functional.mse_loss(temp.view(-1), adaptation_labels)
-        learn.adapt(train_error)
+        # cluster_loss = cl_loss(c)
+        # total_loss = train_error + config['cluster_loss_weight'] * cluster_loss
+        total_loss = train_error + config['kmeans_loss_weight'] * k_loss
+        learn.adapt(total_loss)
+        if is_print:
+            # print("in support:\t", round(k_loss.item(),4))
+            pass
 
-    predictions = learn(evaluation_data, None, training=False, adaptation_data=adaptation_data,
-                        adaptation_labels=adaptation_labels)
+    predictions, c, k_loss = learn(evaluation_data, None, training=False, adaptation_data=adaptation_data,
+                                   adaptation_labels=adaptation_labels)
     valid_error = torch.nn.functional.mse_loss(predictions.view(-1), evaluation_labels)
+    # cluster_loss = cl_loss(c)
+    # total_loss = valid_error + config['cluster_loss_weight'] * cluster_loss
+    total_loss = valid_error + config['kmeans_loss_weight'] * k_loss
+
+    if is_print:
+
+        # print("in query:\t", round(k_loss.item(),4))
+        print(c[0].detach().cpu().numpy(),"\t",round(k_loss.item(),3),"\n")
+
+    # if random.random() < 0.05:
+    #     print("cl:", round(cluster_loss.item()), "\t c:", c[0].cpu().data.numpy())
 
     if get_predictions:
-        return valid_error, predictions
-    return valid_error
+        return total_loss, predictions
+    return total_loss, c, k_loss.item()

learnToLearn.py

 from torch.nn import functional as F
 
 
+def data_batching(indexes, C_distribs, batch_size, training_set_size, num_clusters):
+    probs = np.squeeze(C_distribs)
+    cs = [np.random.choice(num_clusters, p=i) for i in probs]
+    num_batch = int(training_set_size / batch_size)
+    res = [[] for i in range(num_batch)]
+    clas = [[] for i in range(num_clusters)]
+
+    for idx, c in zip(indexes, cs):
+        clas[c].append(idx)
+
+    t = np.array([len(i) for i in clas])
+    t = t / t.sum()
+
+    dif = list(set(list(np.arange(training_set_size))) - set(indexes[0:(num_batch * batch_size)]))
+    cnt = 0
+
+    for i in range(len(res)):
+        for j in range(batch_size):
+            temp = np.random.choice(num_clusters, p=t)
+            if len(clas[temp]) > 0:
+                res[i].append(clas[temp].pop(0))
+            else:
+                # res[i].append(indexes[training_set_size-1-cnt])
+                res[i].append(random.choice(dif))
+                cnt = cnt + 1
+
+    res = np.random.permutation(res)
+    final_result = np.array(res).flatten()
+    return final_result
+
+
 def parse_args():
     print("==============")
     parser = argparse.ArgumentParser([], description='Fast Context Adaptation via Meta-Learning (CAVIA),'
     print("==============\n")
 
     parser.add_argument('--seed', type=int, default=53)
-    # parser.add_argument('--task', type=str, default='multi', help='problem setting: sine or celeba')
-    # parser.add_argument('--tasks_per_metaupdate', type=int, default=32,
-    #                     help='number of tasks in each batch per meta-update')
-    #
-    # parser.add_argument('--lr_inner', type=float, default=5e-6, help='inner-loop learning rate (per task)')
-    # parser.add_argument('--lr_meta', type=float, default=5e-5,
-    #                     help='outer-loop learning rate (used with Adam optimiser)')
+    parser.add_argument('--task', type=str, default='multi', help='problem setting: sine or celeba')
+    parser.add_argument('--tasks_per_metaupdate', type=int, default=32,
+                        help='number of tasks in each batch per meta-update')
+
+    parser.add_argument('--lr_inner', type=float, default=5e-6, help='inner-loop learning rate (per task)')
+    parser.add_argument('--lr_meta', type=float, default=5e-5,
+                        help='outer-loop learning rate (used with Adam optimiser)')
     # parser.add_argument('--lr_meta_decay', type=float, default=0.9, help='decay factor for meta learning rate')
-    #
-    # parser.add_argument('--inner', type=int, default=1,
-    #                     help='number of gradient steps in inner loop (during training)')
-    # parser.add_argument('--inner_eval', type=int, default=1,
-    #                     help='number of gradient updates at test time (for evaluation)')
+
+    parser.add_argument('--inner', type=int, default=1,
+                        help='number of gradient steps in inner loop (during training)')
+    parser.add_argument('--inner_eval', type=int, default=1,
+                        help='number of gradient updates at test time (for evaluation)')
 
     parser.add_argument('--first_order', action='store_true', default=False,
                         help='run first order approximation of CAVIA')
                         help='run adaptation transform')
     parser.add_argument('--transformer', type=str, default="kronoker",
                         help='transformer type')
-    parser.add_argument('--meta_algo', type=str, default="metasgd",
+    parser.add_argument('--meta_algo', type=str, default="gbml",
                         help='MAML/MetaSGD/GBML')
     parser.add_argument('--gpu', type=int, default=0,
                         help='number of gpu to run the code')
     return args
 
 
+from torch.nn import functional as F
+
+
+def kl_loss(C_distribs):
+    # batchsize * k
+    C_distribs = torch.stack(C_distribs).squeeze()
+
+    # print("injam:",len(C_distribs))
+    # print(C_distribs[0].shape)
+    # batchsize * k
+    # print("injam2",C_distribs)
+    C_distribs_sq = torch.pow(C_distribs, 2)
+    # print("injam3",C_distribs_sq)
+    # 1*k
+    C_distribs_sum = torch.sum(C_distribs, dim=0, keepdim=True)
+    # print("injam4",C_distribs_sum)
+    # batchsize * k
+    temp = C_distribs_sq / C_distribs_sum
+    # print("injam5",temp)
+    # batchsize * 1
+    temp_sum = torch.sum(temp, dim=1, keepdim=True)
+    # print("injam6",temp_sum)
+    target_distribs = temp / temp_sum
+    # print("injam7",target_distribs)
+    # calculate the kl loss
+    clustering_loss = F.kl_div(C_distribs.log(), target_distribs, reduction='batchmean')
+    # print("injam8",clustering_loss)
+    return clustering_loss
+
+
 if __name__ == '__main__':
     args = parse_args()
     print(args)
     if config['use_cuda']:
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
         os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
-    master_path = "/media/external_10TB/10TB/maheri/new_data_dir3"
+    master_path = "/media/external_10TB/10TB/maheri/define_task_melu_data2"
     config['master_path'] = master_path
 
     # DATA GENERATION
     fc2_out_dim = config['second_fc_hidden_dim']
     use_cuda = config['use_cuda']
 
-    fc1 = torch.nn.Linear(fc1_in_dim, fc2_in_dim)
-    fc2 = torch.nn.Linear(fc2_in_dim, fc2_out_dim)
-    linear_out = torch.nn.Linear(fc2_out_dim, 1)
-    head = torch.nn.Sequential(fc1, fc2, linear_out)
+    # fc1 = torch.nn.Linear(fc1_in_dim, fc2_in_dim)
+    # fc2 = torch.nn.Linear(fc2_in_dim, fc2_out_dim)
+    # linear_out = torch.nn.Linear(fc2_out_dim, 1)
+    # head = torch.nn.Sequential(fc1, fc2, linear_out)
 
     if use_cuda:
         emb = EmbeddingModule(config).cuda()
         transform = l2l.optim.ModuleTransform(torch.nn.Linear)
 
     trainer = Trainer(config)
+    tr = trainer
 
     # define meta algorithm
     if args.meta_algo == "maml":
         trainer = l2l.algorithms.MAML(trainer, lr=args.lr_inner, first_order=args.first_order)
     elif args.meta_algo == 'metasgd':
-        trainer = l2l.algorithms.MetaSGD(trainer, lr=config['local_lr'], first_order=args.first_order)
+        trainer = l2l.algorithms.MetaSGD(trainer, lr=args.lr_inner, first_order=args.first_order)
     elif args.meta_algo == 'gbml':
-        trainer = l2l.algorithms.GBML(trainer, transform=transform, lr=config['local_lr'],
+        trainer = l2l.algorithms.GBML(trainer, transform=transform, lr=args.lr_inner,
                                       adapt_transform=args.adapt_transform,
                                       first_order=args.first_order)
 
 
     print("\n\n\n")
 
-    for iteration in range(args.epochs):
+    for iteration in range(config['num_epoch']):
+
+        if iteration == 0:
+            print("changing cluster centroids started ...")
+            indexes = list(np.arange(training_set_size))
+            supp_xs, supp_ys, query_xs, query_ys = [], [], [], []
+            for idx in range(0, 2500):
+                supp_xs.append(pickle.load(open("{}/warm_state/supp_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
+                supp_ys.append(pickle.load(open("{}/warm_state/supp_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
+                query_xs.append(
+                    pickle.load(open("{}/warm_state/query_x_{}.pkl".format(master_path, indexes[idx]), "rb")))
+                query_ys.append(
+                    pickle.load(open("{}/warm_state/query_y_{}.pkl".format(master_path, indexes[idx]), "rb")))
+            batch_sz = len(supp_xs)
+
+            user_embeddings = []
+
+            for task in range(batch_sz):
+                # Compute meta-training loss
+                supp_xs[task] = supp_xs[task].cuda()
+                supp_ys[task] = supp_ys[task].cuda()
+                # query_xs[task] = query_xs[task].cuda()
+                # query_ys[task] = query_ys[task].cuda()
+                temp_sxs = emb(supp_xs[task])
+                # temp_qxs = emb(query_xs[task])
+                y = supp_ys[task].view(-1, 1)
+                input_pairs = torch.cat((temp_sxs, y), dim=1)
+                task_embed = tr.cluster_module.input_to_hidden(input_pairs)
+
+                # todo : may be useless
+                mean_task = tr.cluster_module.aggregate(task_embed)
+                user_embeddings.append(mean_task.detach().cpu().numpy())
+
+                supp_xs[task] = supp_xs[task].cpu()
+                supp_ys[task] = supp_ys[task].cpu()
 
-        num_batch = int(training_set_size / batch_size)
-        indexes = list(np.arange(training_set_size))
-        random.shuffle(indexes)
+            from sklearn.cluster import KMeans
+
+            user_embeddings = np.array(user_embeddings)
+            kmeans_model = KMeans(n_clusters=config['cluster_k'], init="k-means++").fit(user_embeddings)
+            tr.cluster_module.array.data = torch.Tensor(kmeans_model.cluster_centers_).cuda()
+
+        if iteration > 0:
+            # indexes = data_batching(indexes, C_distribs, batch_size, training_set_size, config['cluster_k'])
+            # random.shuffle(indexes)
+            C_distribs = []
+        else:
+            num_batch = int(training_set_size / batch_size)
+            indexes = list(np.arange(training_set_size))
+            random.shuffle(indexes)
 
         for i in range(num_batch):
             meta_train_error = 0.0
+            meta_cluster_error = 0.0
             optimizer.zero_grad()
             print("EPOCH: ", iteration, " BATCH: ", i)
             supp_xs, supp_ys, query_xs, query_ys = [], [], [], []
                     query_xs[j] = query_xs[j].cuda()
                     query_ys[j] = query_ys[j].cuda()
 
+            C_distribs = []
             for task in range(batch_sz):
                 # Compute meta-training loss
                 # sxs = supp_xs[task].cuda()
                 temp_sxs = emb(supp_xs[task])
                 temp_qxs = emb(query_xs[task])
 
-                evaluation_error = fast_adapt(learner,
-                                              temp_sxs,
-                                              temp_qxs,
-                                              supp_ys[task],
-                                              query_ys[task],
-                                              config['inner'])
+                evaluation_error, c, k_loss = fast_adapt(learner,
+                                                         temp_sxs,
+                                                         temp_qxs,
+                                                         supp_ys[task],
+                                                         query_ys[task],
+                                                         config['inner'],
+                                                         epoch=iteration)
 
-                evaluation_error.backward()
+                # C_distribs.append(c)
+                evaluation_error.backward(retain_graph=True)
                 meta_train_error += evaluation_error.item()
+                meta_cluster_error += k_loss
 
                 # supp_xs[task].cpu()
                 # query_xs[task].cpu()
             # Print some metrics
             print('Iteration', iteration)
             print('Meta Train Error', meta_train_error / batch_sz)
+            print('KL Train Error', meta_cluster_error / batch_sz)
+
+            # clustering_loss = config['kl_loss_weight'] * kl_loss(C_distribs)
+            # clustering_loss.backward()
+            # print("kl_loss:", round(clustering_loss.item(), 8), "\t", C_distribs[0].cpu().detach().numpy())
+
+            # if i != (num_batch - 1):
+            #     C_distribs = []
 
             # Average the accumulated gradients and optimize
             for p in all_parameters:
             gc.collect()
             print("===============================================\n")
 
-        if iteration % 2 == 0 or iteration>0:
-            # testing
-            print("start of test phase")
-            trainer.eval()
-
-            with open("results2.txt", "a") as f:
-                f.write("epoch:{}\n".format(iteration))
-
-            for test_state in ['user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
-                test_dataset = None
-                test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
-                supp_xs_s = []
-                supp_ys_s = []
-                query_xs_s = []
-                query_ys_s = []
-                gc.collect()
-
-                print("===================== " + test_state + " =====================")
-                mse, ndc1, ndc3 = test(emb, trainer, test_dataset, batch_size=config['batch_size'],num_epoch=config['num_epoch'],test_state=test_state,args=args)
-                with open("results2.txt", "a") as f:
-                    f.write("{}\t{}\t{}\n".format(mse, ndc1, ndc3))
-                print("===================================================")
-                del (test_dataset)
-                gc.collect()
-
-            trainer.train()
-            with open("results2.txt", "a") as f:
-                f.write("\n")
-            print("\n\n\n")
+        # if iteration % 2 == 0 and iteration != 0:
+        #     # testing
+        #     print("start of test phase")
+        #     trainer.eval()
+        #
+        #     with open("results2.txt", "a") as f:
+        #         f.write("epoch:{}\n".format(iteration))
+        #
+        #     for test_state in ['user_cold_state', 'item_cold_state', 'user_and_item_cold_state']:
+        #         test_dataset = None
+        #         test_set_size = int(len(os.listdir("{}/{}".format(master_path, test_state))) / 4)
+        #         supp_xs_s = []
+        #         supp_ys_s = []
+        #         query_xs_s = []
+        #         query_ys_s = []
+        #         gc.collect()
+        #
+        #         print("===================== " + test_state + " =====================")
+        #         mse, ndc1, ndc3 = test(emb, trainer, test_dataset, batch_size=config['batch_size'],
+        #                                num_epoch=config['num_epoch'], test_state=test_state, args=args)
+        #         with open("results2.txt", "a") as f:
+        #             f.write("{}\t{}\t{}\n".format(mse, ndc1, ndc3))
+        #         print("===================================================")
+        #         del (test_dataset)
+        #         gc.collect()
+        #
+        #     trainer.train()
+        #     with open("results2.txt", "a") as f:
+        #         f.write("\n")
+        #     print("\n\n\n")
 
     # save model
     # final_model = torch.nn.Sequential(emb, head)