triplet_loss.py

Go to the documentation of this file.

# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""
 
import torch
import torch.nn.functional as F
 
from fastreid.utils import comm
from fastreid.layers import GatherLayer
from .utils import concat_all_gather, euclidean_dist, normalize
 
 
def softmax_weights(dist, mask):
    max_v = torch.max(dist * mask, dim=1, keepdim=True)[0]
    diff = dist - max_v
    Z = torch.sum(torch.exp(diff) * mask, dim=1, keepdim=True) + 1e-6  # avoid division by zero
    W = torch.exp(diff) * mask / Z
    return W
 
 
def hard_example_mining(dist_mat, is_pos, is_neg):
    """For each anchor, find the hardest positive and negative sample.
    Args:
      dist_mat: pair wise distance between samples, shape [N, M]
      is_pos: positive index with shape [N, M]
      is_neg: negative index with shape [N, M]
    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
    N = dist_mat.size(0)
    # print(f'N = {N}')
    # print(f'dist_mat = {dist_mat}')
    # print(f'dist_mat.shape = {dist_mat.shape}')
    # print(f'is_pos = {is_pos}')
    # print(f'is_pos.shape = {is_pos.shape}')
    # print(f'is_pos_sum = {torch.sum(is_pos)}')
    # print(f'dist_mat[is_pos] = {dist_mat[is_pos]}')
    # print(f'dist_mat[is_pos].shape = {dist_mat[is_pos].shape}')
    # print(f'dist_mat[is_pos].contiguouse() = {dist_mat[is_pos].contiguous()}')
    # print(f'dist_mat[is_pos].contiguouse().shape = {dist_mat[is_pos].contiguous().shape}')
    # print(f'dist_mat[is_pos].contiguouse().view(N,-1) = {dist_mat[is_pos].contiguous().view(N,-1)}')
 
 
    # `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].reshape(N,-1), 1, keepdim=True)
        # 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].reshape(N,-1), 1, keepdim=True)
        # 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)
 
    return dist_ap, dist_an
 
 
def weighted_example_mining(dist_mat, is_pos, is_neg):
    """For each anchor, find the weighted positive and negative sample.
    Args:
      dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N]
      is_pos:
      is_neg:
    Returns:
      dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
      dist_an: pytorch Variable, distance(anchor, negative); shape [N]
    """
    assert len(dist_mat.size()) == 2
 
    is_pos = is_pos.float()
    is_neg = is_neg.float()
    dist_ap = dist_mat * is_pos
    dist_an = dist_mat * is_neg
 
    weights_ap = softmax_weights(dist_ap, is_pos)
    weights_an = softmax_weights(-dist_an, is_neg)
 
    dist_ap = torch.sum(dist_ap * weights_ap, dim=1)
    dist_an = torch.sum(dist_an * weights_an, dim=1)
 
    return dist_ap, dist_an
 
 
def triplet_loss(embedding, targets, margin, norm_feat, hard_mining):
    r"""Modified from Tong Xiao's open-reid (https://github.com/Cysu/open-reid).
    Related Triplet Loss theory can be found in paper 'In Defense of the Triplet
    Loss for Person Re-Identification'."""
 
    if norm_feat: embedding = normalize(embedding, axis=-1)
 
    # For distributed training, gather all features from different process.
    if comm.get_world_size() > 1:
        all_embedding = torch.cat(GatherLayer.apply(embedding), dim=0)
        all_targets = concat_all_gather(targets)
    else:
        all_embedding = embedding
        all_targets = targets
 
    dist_mat = euclidean_dist(embedding, all_embedding)
 
    N, M = dist_mat.size()
    is_pos = targets.view(N, 1).expand(N, M).eq(all_targets.view(M, 1).expand(M, N).t())
    is_neg = targets.view(N, 1).expand(N, M).ne(all_targets.view(M, 1).expand(M, N).t())
 
    if hard_mining:
        dist_ap, dist_an = hard_example_mining(dist_mat, is_pos, is_neg)
    else:
        dist_ap, dist_an = weighted_example_mining(dist_mat, is_pos, is_neg)
 
    y = dist_an.new().resize_as_(dist_an).fill_(1)
 
    if margin > 0:
        loss = F.margin_ranking_loss(dist_an, dist_ap, y, margin=margin)
    else:
        loss = F.soft_margin_loss(dist_an - dist_ap, y)
        # fmt: off
        if loss == float('Inf'): loss = F.margin_ranking_loss(dist_an, dist_ap, y, margin=0.3)
        # fmt: on
 
    return loss