import paddle
import paddle.nn.functional as F

from ..registry import LOSSES
from .base import BaseWeightedLoss

@LOSSES.register()
class TripletLoss(BaseWeightedLoss):
    """Triplet loss with hard positive/negative mining.

    Reference:
    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.

    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.

    Args:
        margin (float): margin for triplet.
    """
    def __init__(self, margin=0.3, distance='euclidean', use_gpu=True):
        super(TripletLoss, self).__init__()
        if distance not in ['euclidean', 'consine']:
            raise KeyError("Unsupported distance: {}".format(distance))
        self.distance = distance
        self.margin = margin
        self.use_gpu = use_gpu
        self.ranking_loss = paddle.nn.MarginRankingLoss(margin=margin,reduction='mean')
        
    def forward(self, inputs, targets):
        """
        Args:
            inputs: feature matrix with shape (batch_size, feat_dim)
            targets: ground truth labels with shape (num_classes)
        """
        n = inputs.shape[0]
        
        # Compute pairwise distance, replace by the official when merged
        if self.distance == 'euclidean':
            dist = paddle.pow(inputs, 2).sum(axis=1, keepdim=True).expand([n, n])
            dist = dist + dist.t()
            dist = dist - 2 * paddle.mm(inputs, inputs.t())
           
            dist = paddle.clip(dist, min=1e-12).sqrt()  # for numerical stability
        elif self.distance == 'consine':
            fnorm = paddle.norm(inputs, p=2, axis=1, keepdim=True)
            l2norm = inputs.div(fnorm.expand_as(inputs))
            dist = - paddle.mm(l2norm, l2norm.t())
        if self.use_gpu: targets = targets.cuda()
        # For each anchor, find the hardest positive and negative
        mask = targets.expand([n, n]).equal(targets.expand([n, n]).t()).astype(int)
        dist_ap, dist_an = [], []
       
        
        for i in range(n):
            dist_ap.append(paddle.masked_select(dist[i], mask[i]==1).max().unsqueeze(0))
            dist_an.append(paddle.masked_select(dist[i], mask[i]==0).min().unsqueeze(0))
        dist_ap = paddle.concat(dist_ap)
        dist_an = paddle.concat(dist_an)
        
        # Compute ranking hinge loss
        y = paddle.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)
        return loss
import paddle
import paddle.nn.functional as F

from ..registry import LOSSES
from .base import BaseWeightedLoss

@LOSSES.register()
class TripletLoss(BaseWeightedLoss):
    """Triplet loss with hard positive/negative mining.

    Reference:
    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.

    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.

    Args:
        margin (float): margin for triplet.
    """
    def __init__(self, margin=0.3, distance='euclidean', use_gpu=True):
        super(TripletLoss, self).__init__()
        if distance not in ['euclidean', 'consine']:
            raise KeyError("Unsupported distance: {}".format(distance))
        self.distance = distance
        self.margin = margin
        self.use_gpu = use_gpu
        self.ranking_loss = paddle.nn.MarginRankingLoss(margin=margin,reduction='mean')
        
    def forward(self, inputs, targets):
        """
        Args:
            inputs: feature matrix with shape (batch_size, feat_dim)
            targets: ground truth labels with shape (num_classes)
        """
        n = inputs.shape[0]
        
        # Compute pairwise distance, replace by the official when merged
        if self.distance == 'euclidean':
            dist = paddle.pow(inputs, 2).sum(axis=1, keepdim=True).expand([n, n])
            dist = dist + dist.t()
            dist = dist - 2 * paddle.mm(inputs, inputs.t())
           
            dist = paddle.clip(dist, min=1e-12).sqrt()  # for numerical stability
        elif self.distance == 'consine':
            fnorm = paddle.norm(inputs, p=2, axis=1, keepdim=True)
            l2norm = inputs.div(fnorm.expand_as(inputs))
            dist = - paddle.mm(l2norm, l2norm.t())
        if self.use_gpu: targets = targets.cuda()
        # For each anchor, find the hardest positive and negative
        mask = targets.expand([n, n]).equal(targets.expand([n, n]).t()).astype(int)
        dist_ap, dist_an = [], []
       
        
        for i in range(n):
            dist_ap.append(paddle.masked_select(dist[i], mask[i]==1).max().unsqueeze(0))
            dist_an.append(paddle.masked_select(dist[i], mask[i]==0).min().unsqueeze(0))
        dist_ap = paddle.concat(dist_ap)
        dist_an = paddle.concat(dist_an)
        
        # Compute ranking hinge loss
        y = paddle.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)
        return loss
Triplet loss with hard positive/negative mining.      Reference:     Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.      Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.      Args:         margin (float): margin for triplet. Args:             inputs: feature matrix with shape (batch_size, feat_dim)             targets: ground truth labels with shape (num_classes) # Compute pairwise distance, replace by the official when merged # for numerical stability # For each anchor, find the hardest positive and negative # Compute ranking hinge loss