# AUTOGENERATED! DO NOT EDIT! File to edit: 08_contrastive_loss.ipynb (unless otherwise specified).

__all__ = ['TripletLoss', 'ContrastiveLoss', 'CosineContrastiveLoss', 'batched_labels', 'XentOldContrastiveLoss',
           'XentLoss', 'XentContrastiveLoss', 'XentContrastiveLoss2', 'BatchContrastiveLoss']

# Cell
from .imports import *

# Cell
class TripletLoss(nn.Module):
    """
    Triplet loss
    Takes embeddings of an anchor sample, a positive sample and a negative sample
    """

    def __init__(self, margin=0.5):
        super(TripletLoss, self).__init__()
        self.margin = margin

    def forward(self, anchor, positive, negative=None, size_average=True):
        if negative is None: negative = positive.flip(dims=[0])
        distance_positive = (anchor - positive).pow(2).sum(1)  # .pow(.5)
        distance_negative = (anchor - negative).pow(2).sum(1)  # .pow(.5)
        losses = F.relu(distance_positive - distance_negative + self.margin)
        return losses.mean() if size_average else losses.sum()


# Cell
# https://github.com/adambielski/siamese-triplet/blob/master/losses.py
class ContrastiveLoss(nn.Module):
    """
    Contrastive loss
    Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise
    """
    def __init__(self, margin=2.0):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin
        self.eps = 1e-9

    def forward(self, output1, output2, target, size_average=True):
        distances = (output2 - output1).pow(2).sum(-1)  # squared distances
#         distances = F.pairwise_distance(output1, output2, keepdim=True).pow(2).sum(-1)  # squared distances
        losses = (target.float() * distances +
                        (1 + -1 * target).float() * F.relu(self.margin - (distances + self.eps).sqrt()).pow(2))
        return losses.mean() if size_average else losses.sum()

# Cell
# https://stackoverflow.com/questions/47107589/how-do-i-use-a-bytetensor-in-a-contrastive-cosine-loss-function
class CosineContrastiveLoss(nn.Module):
    def __init__(self, margin=2.0):
        super().__init__()
        self.margin = margin

    def forward(self, output1, output2, label):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)
        loss_cos_con = torch.mean((label) * torch.div(torch.pow((1.0-cos_sim), 2), 4) +
                                    (1-label) * torch.pow(cos_sim * torch.lt(cos_sim, self.margin), 2))
        return loss_cos_con

# Cell
def batched_labels(output1, output2, onehot=True):
    bs = output1.shape[0]
    rp = [1]*len(output1.shape)
    o1 = output1.repeat(*rp,bs).view(bs,*output1.shape)
    labels = torch.arange(bs, device=output1.device)
    if onehot: labels = torch.eye(o1.shape[0], device=output1.device)[labels]
    return o1, output2, labels

# Cell
# https://arxiv.org/pdf/2002.05709.pdf
class XentOldContrastiveLoss(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        xent_loss = F.cross_entropy(cos_sim,labels.long())
        return xent_loss

# Cell
from pytorch_metric_learning import losses
class XentLoss(losses.NTXentLoss):
    def forward(self, output1, output2):
        stacked = torch.cat((output1, output2), dim=0)
        labels = torch.arange(output1.shape[0]).repeat(2)
        return super().forward(stacked, labels, None)

# Cell
# https://arxiv.org/pdf/2002.05709.pdf
class XentContrastiveLoss(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        cexp = torch.exp(cos_sim)
        neg_denom = (cexp*(1-labels)).sum(dim=-1)
        lsoft = torch.log(cexp/neg_denom)
        lsoft = torch.sum(-labels * lsoft, dim=-1)
        print(lsoft)
        return lsoft.mean()

# Cell
class XentContrastiveLoss2(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        cexp = torch.exp(cos_sim)
        x = (cexp * labels).sum(dim=-1)
        denom = cexp.sum(dim=-1) - x
        lsoft = -torch.log(x/denom)
        print(lsoft)
        return lsoft.mean()

# Cell
class BatchContrastiveLoss(nn.Module):
    def __init__(self, loss_func):
        super().__init__()
        self.loss_func = loss_func
        self.onehot = not isinstance(loss_func, XentOldContrastiveLoss)

    def forward(self, output1, output2):
        output1, output2, labels = batched_labels(output1, output2, self.onehot)
        return self.loss_func(output1, output2, labels)
# AUTOGENERATED! DO NOT EDIT! File to edit: 08_contrastive_loss.ipynb (unless otherwise specified).

__all__ = ['TripletLoss', 'ContrastiveLoss', 'CosineContrastiveLoss', 'batched_labels', 'XentOldContrastiveLoss',
           'XentLoss', 'XentContrastiveLoss', 'XentContrastiveLoss2', 'BatchContrastiveLoss']

# Cell
from .imports import *

# Cell
class TripletLoss(nn.Module):
    """
    Triplet loss
    Takes embeddings of an anchor sample, a positive sample and a negative sample
    """

    def __init__(self, margin=0.5):
        super(TripletLoss, self).__init__()
        self.margin = margin

    def forward(self, anchor, positive, negative=None, size_average=True):
        if negative is None: negative = positive.flip(dims=[0])
        distance_positive = (anchor - positive).pow(2).sum(1)  # .pow(.5)
        distance_negative = (anchor - negative).pow(2).sum(1)  # .pow(.5)
        losses = F.relu(distance_positive - distance_negative + self.margin)
        return losses.mean() if size_average else losses.sum()


# Cell
# https://github.com/adambielski/siamese-triplet/blob/master/losses.py
class ContrastiveLoss(nn.Module):
    """
    Contrastive loss
    Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise
    """
    def __init__(self, margin=2.0):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin
        self.eps = 1e-9

    def forward(self, output1, output2, target, size_average=True):
        distances = (output2 - output1).pow(2).sum(-1)  # squared distances
#         distances = F.pairwise_distance(output1, output2, keepdim=True).pow(2).sum(-1)  # squared distances
        losses = (target.float() * distances +
                        (1 + -1 * target).float() * F.relu(self.margin - (distances + self.eps).sqrt()).pow(2))
        return losses.mean() if size_average else losses.sum()

# Cell
# https://stackoverflow.com/questions/47107589/how-do-i-use-a-bytetensor-in-a-contrastive-cosine-loss-function
class CosineContrastiveLoss(nn.Module):
    def __init__(self, margin=2.0):
        super().__init__()
        self.margin = margin

    def forward(self, output1, output2, label):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)
        loss_cos_con = torch.mean((label) * torch.div(torch.pow((1.0-cos_sim), 2), 4) +
                                    (1-label) * torch.pow(cos_sim * torch.lt(cos_sim, self.margin), 2))
        return loss_cos_con

# Cell
def batched_labels(output1, output2, onehot=True):
    bs = output1.shape[0]
    rp = [1]*len(output1.shape)
    o1 = output1.repeat(*rp,bs).view(bs,*output1.shape)
    labels = torch.arange(bs, device=output1.device)
    if onehot: labels = torch.eye(o1.shape[0], device=output1.device)[labels]
    return o1, output2, labels

# Cell
# https://arxiv.org/pdf/2002.05709.pdf
class XentOldContrastiveLoss(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        xent_loss = F.cross_entropy(cos_sim,labels.long())
        return xent_loss

# Cell
from pytorch_metric_learning import losses
class XentLoss(losses.NTXentLoss):
    def forward(self, output1, output2):
        stacked = torch.cat((output1, output2), dim=0)
        labels = torch.arange(output1.shape[0]).repeat(2)
        return super().forward(stacked, labels, None)

# Cell
# https://arxiv.org/pdf/2002.05709.pdf
class XentContrastiveLoss(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        cexp = torch.exp(cos_sim)
        neg_denom = (cexp*(1-labels)).sum(dim=-1)
        lsoft = torch.log(cexp/neg_denom)
        lsoft = torch.sum(-labels * lsoft, dim=-1)
        print(lsoft)
        return lsoft.mean()

# Cell
class XentContrastiveLoss2(nn.Module):
    def __init__(self, temp=0.5):
        super().__init__()
        self.temp = temp

    def forward(self, output1, output2, labels):
        cos_sim = F.cosine_similarity(output1, output2, dim=-1)/self.temp
        cexp = torch.exp(cos_sim)
        x = (cexp * labels).sum(dim=-1)
        denom = cexp.sum(dim=-1) - x
        lsoft = -torch.log(x/denom)
        print(lsoft)
        return lsoft.mean()

# Cell
class BatchContrastiveLoss(nn.Module):
    def __init__(self, loss_func):
        super().__init__()
        self.loss_func = loss_func
        self.onehot = not isinstance(loss_func, XentOldContrastiveLoss)

    def forward(self, output1, output2):
        output1, output2, labels = batched_labels(output1, output2, self.onehot)
        return self.loss_func(output1, output2, labels)
# AUTOGENERATED! DO NOT EDIT! File to edit: 08_contrastive_loss.ipynb (unless otherwise specified). # Cell # Cell Triplet loss     Takes embeddings of an anchor sample, a positive sample and a negative sample # .pow(.5) # .pow(.5) # Cell # https://github.com/adambielski/siamese-triplet/blob/master/losses.py Contrastive loss     Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise # squared distances #         distances = F.pairwise_distance(output1, output2, keepdim=True).pow(2).sum(-1)  # squared distances # Cell # https://stackoverflow.com/questions/47107589/how-do-i-use-a-bytetensor-in-a-contrastive-cosine-loss-function # Cell # Cell # https://arxiv.org/pdf/2002.05709.pdf # Cell # Cell # https://arxiv.org/pdf/2002.05709.pdf # Cell # Cell