mini-batch的TripleLoss實現(Pytorch)

以前都是直接調用別人的， 但是詳細實現沒有了解過， 今天自己實現一把。簡單來說， 找出batch中每個anchor對應的最大正樣本postive和最小負樣本nagetive，然后距離max(a-p)和min(a-n)做差即可。

class TripleLoss(nn.Module):
    def __init__(self, margin=0.3):
        super(TripleLoss, self).__init__()
        self.margin = margin # 閾值
        self.rank_loss = nn.MarginRankingLoss(margin=margin)

    def forward(self, inputs, labels, norm=False):      
        dist_mat = self.euclidean_dist(inputs, inputs, norm=norm)  # 距離矩陣
        dist_ap, dist_an = self.hard_sample(dist_mat, labels) # 取出每個anchor對應的最大
        y = torch.ones_like(dist_an) # 系數矩陣，1/-1
        loss = self.rank_loss(dist_ap, dist_an, y)
        return loss

    @staticmethod
    def hard_sample( dist_mat, labels, ):
        # 距離矩陣的尺寸是 (batch_size, batch_size)
        assert len(dist_mat.size()) == 2
        assert dist_mat.size(0) == dist_mat.size(1)
        N = dist_mat.size(0)

        # 選出所有正負樣本對
        is_pos = labels.expand(N, N).eq(labels.expand(N, N).t()) # 兩兩組合， 取label相同的a-p
        is_neg = labels.expand(N, N).ne(labels.expand(N, N).t()) # 兩兩組合， 取label不同的a-n

        list_ap, list_an = [], []
        # 取出所有正樣本對和負樣本對的距離值
        for i in range(N):
            list_ap.append( dist_mat[i][is_pos[i]].max().unsqueeze(0) ) 
            list_an.append( dist_mat[i][is_neg[i]].max().unsqueeze(0) )
            dist_ap = torch.cat(list_ap)  # 將list里的tensor拼接成新的tensor
            dist_an = torch.cat(list_an)
        return dist_ap, dist_an

    @staticmethod
    def normalize(x, axis=1):
        x = 1.0*x / (torch.norm(x, 2, axis, keepdim=True) + 1e-12)
        return x

    @staticmethod
    def euclidean_dist(x, y, norm=True):
        if norm:
            x = self.normalize(x)
            y = self.normalize(y)
        m, n = x.size(0), y.size(0)
        xx = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n)
        yy = torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t()
        dist = xx + yy # 任意的兩個樣本組合， 求第二范數后求和 x^2 + y^2
        dist.addmm_( 1, -2, x, y.t() ) # (x-y)^2 = x^2 + y^2 - 2xy
        dist = dist.clamp(min=1e-12).sqrt()
        return dist