|
- from __future__ import absolute_import
-
- import torch
- from torch import nn
-
-
- def normalize(x, axis=-1):
- """Normalizing to unit length along the specified dimension.
- Args:
- x: pytorch Variable
- Returns:
- x: pytorch Variable, same shape as input
- """
- x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
- return x
-
-
- def euclidean_dist(x, y):
- """
- Args:
- x: pytorch Variable, with shape [m, d]
- y: pytorch Variable, with shape [n, d]
- Returns:
- dist: pytorch Variable, with shape [m, n]
- """
- m, n = x.size(0), y.size(0)
- xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
- yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
- dist = xx + yy
- # dist.addmm_(1, -2, x, y.t())
- dist.addmm_( x, y.t(), beta = 1, alpha = -2)
- dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
- return dist
-
-
- def hard_example_mining(dist_mat, labels, return_inds=False):
- """For each anchor, find the hardest positive and negative sample.
- Args:
- dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N]
- labels: pytorch LongTensor, with shape [N]
- return_inds: whether to return the indices. Save time if `False`(?)
- Returns:
- dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
- dist_an: pytorch Variable, distance(anchor, negative); shape [N]
- p_inds: pytorch LongTensor, with shape [N];
- indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
- n_inds: pytorch LongTensor, with shape [N];
- indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1
- NOTE: Only consider the case in which all labels have same num of samples,
- thus we can cope with all anchors in parallel.
- """
-
- assert len(dist_mat.size()) == 2
- assert dist_mat.size(0) == dist_mat.size(1)
- N = dist_mat.size(0)
-
- # shape [N, N]
- is_pos = labels.expand(N, N).eq(labels.expand(N, N).t())
- is_neg = labels.expand(N, N).ne(labels.expand(N, N).t())
-
- '''
- # `dist_ap` means distance(anchor, positive)
- # both `dist_ap` and `relative_p_inds` with shape [N, 1]
-
- dist_ap, relative_p_inds = torch.max(dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True)
- # `dist_an` means distance(anchor, negative)
- # both `dist_an` and `relative_n_inds` with shape [N, 1]
- dist_an, relative_n_inds = torch.min(dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True)
- # shape [N]
- dist_ap = dist_ap.squeeze(1)
- dist_an = dist_an.squeeze(1)
- '''
-
- dist_ap, relative_p_inds = torch.max(dist_mat.masked_fill(is_neg, 0), 1, keepdim=True)
- dist_an, relative_n_inds = torch.min(dist_mat.masked_fill(is_pos, float("inf")), 1, keepdim=True)
-
- # dist_ap1, dist_an1 = [], []
- # for i in range(N):
- # dist_ap1.append(dist_mat[i][is_pos[i]].max().unsqueeze(0))
- # dist_an1.append(dist_mat[i][is_pos[i] == 0].min().unsqueeze(0))
-
- # dist_ap1, dist_an1 = torch.tensor(dist_ap1), torch.tensor(dist_an1)
-
- if return_inds:
- # shape [N, N]
- ind = (labels.new().resize_as_(labels).copy_(torch.arange(0, N).long()).unsqueeze(0).expand(N, N))
- # shape [N, 1]
- p_inds = torch.gather(ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data)
- n_inds = torch.gather(ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data)
- # shape [N]
- p_inds = p_inds.squeeze(1)
- n_inds = n_inds.squeeze(1)
- return dist_ap, dist_an, p_inds, n_inds
-
- return dist_ap, dist_an
-
-
- class TripletLoss(nn.Module):
- def __init__(self, margin=None):
- super(TripletLoss, self).__init__()
- self.margin = margin
- if self.margin is not None:
- self.ranking_loss = nn.MarginRankingLoss(margin=margin)
- else:
- self.ranking_loss = nn.SoftMarginLoss()
-
- def forward(self, global_feat, labels, normalize_feature=False):
- if normalize_feature:
- global_feat = normalize(global_feat, axis=-1)
- # shape [N, N]
- dist_mat = euclidean_dist(global_feat, global_feat)
- dist_ap, dist_an = hard_example_mining(dist_mat, labels)
- y = dist_an.new().resize_as_(dist_an).fill_(1)
- if self.margin is not None:
- loss = self.ranking_loss(dist_an, dist_ap, y)
- else:
- loss = self.ranking_loss(dist_an - dist_ap, y)
- prec = (dist_an.data > dist_ap.data).sum().float() / y.size(0)
- return loss, prec, dist_ap, dist_an
-
-
- class Center_loss(nn.Module):
- """Center loss.
- Reference:
- Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
- Args:
- num_classes (int): number of classes.
- feat_dim (int): feature dimension.
- """
-
- def __init__(self, num_classes=751, feat_dim=2048, use_gpu=True):
- super(Center_loss, self).__init__()
- self.num_classes = num_classes
- self.feat_dim = feat_dim
- self.use_gpu = use_gpu
-
- if self.use_gpu:
- self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda())
- else:
- self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))
-
- def forward(self, x, labels):
- """
- Args:
- x: feature matrix with shape (batch_size, feat_dim).
- labels: ground truth labels with shape (num_classes).
- """
- assert x.size(0) == labels.size(0), "features.size(0) is not equal to labels.size(0)"
-
- batch_size = x.size(0)
- distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \
- torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t()
- distmat.addmm_(1, -2, x, self.centers.t())
-
- classes = torch.arange(self.num_classes).long()
- if self.use_gpu: classes = classes.cuda()
- labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
- mask = labels.eq(classes.expand(batch_size, self.num_classes))
-
- dist = distmat * mask.float()
- loss = dist.clamp(min=1e-12, max=1e+12).sum() / batch_size
- #dist = []
- #for i in range(batch_size):
- # value = distmat[i][mask[i]]
- # value = value.clamp(min=1e-12, max=1e+12) # for numerical stability
- # dist.append(value)
- #dist = torch.cat(dist)
- #loss = dist.mean()
- return loss
|