some of the investigation types

info_nce.py

@@ -4,6 +4,7 @@
 
 from scale_grad import ScaleGrad
 from mp_utils import AllGather
+import copy
 
 __all__ = ['simclr_loss']
 
@@ -24,7 +25,9 @@ def __init__(self, tau: float = 0.07, grad_scaling=False):
         self.tau = tau
         self.grad_scaling = grad_scaling
 
-    def forward(self, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
+    def forward(self, z1: torch.Tensor, z2: torch.Tensor, epoch) -> torch.Tensor:
+
+        run_type = ["label_smoothing", "top_k_based_on_cos_sim", "scheduled_cos_sim", "scheduled_top_k_cos_sim"][3]
 
         # Collect from all gpu
         z1 = AllGather.apply(z1)
@@ -41,16 +44,89 @@ def forward(self, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
         ones = torch.ones(n // 2).to(z.device)
         labels = ones.diagflat(n // 2) + ones.diagflat(-n // 2)
 
+        if run_type == "label_smoothing":
+            labels = labels.clamp(min=0.01)
+            labels = F.normalize(labels, dim=1)
+
         # Note: The following code might require a large amount of memory
         # in case of large batch size
         sim_m = z @ z.T
+        F.normalize(z, dim=1, p=2)
+        if run_type == "top_k_based_on_cos_sim":
+            # pairwise similarities between embeddings
+            cos_sim_m = F.softmax(z @ z.T, dim=1)
+            # masking  for ruling-out pairs from the same siamese branch
+            branch_mask_zeros = torch.zeros((int(cos_sim_m.shape[0]/2), int(cos_sim_m.shape[1]/2)))
+            branch_mask_ones = torch.ones((int(cos_sim_m.shape[0]/2), int(cos_sim_m.shape[1]/2)))
+            branch_mask = torch.cat(
+                [torch.cat([branch_mask_zeros, branch_mask_ones], dim=1),
+                torch.cat([branch_mask_ones, branch_mask_zeros], dim=1)],
+                dim=0
+            ).to(cos_sim_m.device)
+            cos_sim_m = torch.mul(branch_mask, cos_sim_m)
+            # selecting top-k similarities
+            top_k, top_k_indices = torch.topk(cos_sim_m, k=10, dim=1)
+            top_sim = torch.zeros_like(cos_sim_m).scatter(dim=1, index=top_k_indices, src=torch.full_like(top_k, 0.75))
+            labels = torch.max(labels, top_sim)
+        elif run_type == "scheduled_cos_sim":
+            # pairwise similarities between embeddings
+            cos_sim_m = F.softmax(z @ z.T, dim=1)
+            # masking  for ruling-out pairs from the same siamese branch
+            branch_mask_zeros = torch.zeros((int(cos_sim_m.shape[0]/2), int(cos_sim_m.shape[1]/2)))
+            branch_mask_ones = torch.ones((int(cos_sim_m.shape[0]/2), int(cos_sim_m.shape[1]/2)))
+            branch_mask = torch.cat(
+                [torch.cat([branch_mask_zeros, branch_mask_ones], dim=1),
+                torch.cat([branch_mask_ones, branch_mask_zeros], dim=1)],
+                dim=0
+            ).to(cos_sim_m.device)
+            cos_sim_m = torch.mul(branch_mask, cos_sim_m)
+            # scheduling
+            schedule_start_epoch = 30
+            schedule_length_epoch = 20
+            if epoch >= schedule_start_epoch:
+                schedule = torch.linspace(start=0, end=1, steps=schedule_length_epoch)
+                index_based_on_epoch = min(epoch - schedule_start_epoch, schedule_length_epoch - 1)
+                labels = (1 - schedule[index_based_on_epoch]) * labels + schedule[index_based_on_epoch] * cos_sim_m
+                labels = F.softmax(labels, dim=1)
+        elif run_type == "scheduled_top_k_cos_sim":
+            # hyper-parameters
+            num_of_top_k = 10  # 10
+            schedule_start_epoch = 30  # 30
+            schedule_length_epoch = 20
+            # pairwise similarities between embeddings
+            cos_sim_m = F.softmax(z @ z.T, dim=1)
+            # masking  for ruling-out pairs from the same siamese branch
+            branch_mask_zeros = torch.zeros((int(cos_sim_m.shape[0] / 2), int(cos_sim_m.shape[1] / 2)))
+            branch_mask_ones = torch.ones((int(cos_sim_m.shape[0] / 2), int(cos_sim_m.shape[1] / 2)))
+            branch_mask = torch.cat(
+                [torch.cat([branch_mask_zeros, branch_mask_ones], dim=1),
+                 torch.cat([branch_mask_ones, branch_mask_zeros], dim=1)],
+                dim=0
+            ).to(cos_sim_m.device)
+            cos_sim_m = torch.mul(branch_mask, cos_sim_m)
+            # selecting top-k similarities
+            top_k, top_k_indices = torch.topk(cos_sim_m, k=num_of_top_k, dim=1)
+            top_sim = torch.zeros_like(cos_sim_m).scatter(dim=1, index=top_k_indices, src=torch.full_like(top_k, 0.75))
+            labels_based_on_similarity = torch.max(labels, top_sim)
+            # scheduling
+            if epoch >= schedule_start_epoch:
+                schedule = torch.linspace(start=0, end=1, steps=schedule_length_epoch)
+                index_based_on_epoch = min(epoch - schedule_start_epoch, schedule_length_epoch - 1)
+                labels = (1 - schedule[index_based_on_epoch]) * labels + schedule[index_based_on_epoch] * labels_based_on_similarity
+                labels = F.normalize(labels, dim=1)
 
         # This is a bit of cheat. Instead of removing cells from
         # the matrix where i==j, instead we set it to a very small value
         sim_m = sim_m.fill_diagonal_(-10) / self.tau
 
         # Get probability distribution
-        sim_m = torch.nn.functional.log_softmax(sim_m, dim=1)
+        sim_true_positive = sim_m.clone()
+        sim_auxiliary_positive = sim_m.clone()
+        true_positive_mask = torch.roll(torch.eye(sim_m.shape[0]), int(sim_m.shape[0]/2), dims=0).bool().cuda()
+        sim_auxiliary_positive.masked_fill_(true_positive_mask, (-10 / self.tau))
+        sim_true_positive = torch.nn.functional.log_softmax(sim_true_positive, dim=1)
+        sim_auxiliary_positive = torch.nn.functional.log_softmax(sim_auxiliary_positive, dim=1)
+        sim_m = torch.where(true_positive_mask.cuda(), sim_true_positive, sim_auxiliary_positive)
 
         # Choose values on which we calculate the loss
         loss = -torch.sum(sim_m * labels) / n

main_linearpredhead.py

@@ -41,7 +41,7 @@
     and callable(models.__dict__[name]))
 
 parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
-parser.add_argument("--runname", default="dev", help="Name of run on tensorboard")
+parser.add_argument("--runname", default="schedule_top_10_cos_sim_from30_through20", help="Name of run on tensorboard")
 parser.add_argument('--data', metavar='DIR',
                     help='path to dataset', default='/shared_data/imagenet')
 parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50',
@@ -102,10 +102,10 @@
                     help='Where to log')
 parser.add_argument("--save-frequency", default=5, help="Frequency of checkpoint saving in epochs")
 
-parser.add_argument("--pred-type", default="linear", help="type of the prediction head",
+parser.add_argument("--pred-type", default="original", help="type of the prediction head",
                     choices=["linear", "random_linear", "predefined_linear", "low_rank_linear", "original"])
 
-parser.add_argument("--loss-type", default="cossim", help="type of loss",
+parser.add_argument("--loss-type", default="infonce", help="type of loss",
                     choices=["cossim", "infonce", "infonce_scalegrad", "cossim_scalegrad"])
 
 parser.add_argument('--fix-p-size', default=-1.0, type=float,
@@ -346,7 +346,9 @@ def cosine_similarioty(x1, x2):
 
 def create_loss(type="cossim", fix_p_size=-1):
     # cossim = CosineSimilarity(inverse=True)
-    fix_p_term = lambda x: ((torch.norm(x, p=2, dim=-1) - fix_p_size) ** 2).mean() if fix_p_size > 0 else lambda x: 0.0
+    fix_p_term = (lambda x: ((torch.norm(x, p=2, dim=-1) - fix_p_size) ** 2).mean()) if fix_p_size > 0 \
+        else lambda x: torch.Tensor([0.0]).to(device=x.device)
+
     if "cossim" in type:
         cossim = torch.nn.CosineSimilarity(dim=1)
         def compute_all_loss(p1s, p2s, z1, z2):
@@ -370,12 +372,12 @@ def compute_all_loss(p1s, p2s, z1, z2):
             grad_scaling = True
         else:
             grad_scaling = False
-        info_nce =  InfoNCE(grad_scaling=grad_scaling)
-        def compute_all_loss(p1s, p2s, z1, z2):
+        info_nce = InfoNCE(grad_scaling=grad_scaling)
+        def compute_all_loss(p1s, p2s, z1, z2, epoch):
             losses = {"optim": dict(),
                       "log": dict()}
             for i, ((bn_dim, p1), p2) in enumerate(zip(p1s.items(), p2s.values())):
-                avg_info_nce = (info_nce(p1, z2).mean() + info_nce(p2, z1).mean()) * 0.5
+                avg_info_nce = (info_nce(p1, z2, epoch).mean() + info_nce(p2, z1, epoch).mean()) * 0.5
                 avg_fix_p_term = (fix_p_term(p1) + fix_p_term(p2)) * 0.5
                 losses["optim"][bn_dim] = avg_info_nce + avg_fix_p_term
                 losses["log"][f"fix_p_term_{bn_dim}"] = avg_fix_p_term
@@ -411,7 +413,7 @@ def train(train_loader, model, predictor, criterion, optimizer, epoch, tb_logger
 
         # compute output and loss
         p1, p2, z1, z2 = model(x1=images[0], x2=images[1])
-        loss_dict = criterion(p1, p2, z1, z2)
+        loss_dict = criterion(p1, p2, z1, z2, epoch=epoch)
 
         loss = .0
         for bn_dim, value in loss_dict["optim"].items():