diff --git a/info_nce.py b/info_nce.py index 136f371..c3c40bb 100644 --- a/info_nce.py +++ b/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 diff --git a/main_linearpredhead.py b/main_linearpredhead.py index e5f5335..542a1a0 100755 --- a/main_linearpredhead.py +++ b/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():