Hi Jhuteau!



 jhuteau:

I want an efficient (vectorized) implementation instead of my current for-loop implementation.
…
How can I efficiently implement the loss with a varying number of positive documents for each anchor ?


Easy peasy:
import math

import torch

print (torch.__version__)

torch.manual_seed (20212021)


# code with tensor operations

def hardest_triplet_margin_lossB (preds, targets, margin, hardest_fraction):
    # preds:             the similarity scores, of shape [num_anchors, num-documents]
    # targets:           the "positiveness" of the scores, of the same shape
    # margin:            the margin-loss margin
    # hardest_fraction:  fraction of weakest "positive" and "negative" scores to analyze
    
    n_doc = preds.shape[1]
    
    pos_mask = (targets >= 0.5).int()
    neg_mask = 1 - pos_mask
    max_preds = preds.max (dim = 1)[0]
    min_preds = preds.min (dim = 1)[0]
    
    pos_sort = pos_mask * preds  +  neg_mask * (1 + max_preds).unsqueeze (1)
    pos_sort = pos_sort.sort (dim = 1, descending = True)[0]
    pos_k = (hardest_fraction * pos_mask.sum (dim = 1)).ceil().maximum (torch.ones (1))
    pos_hard_mask = ((n_doc - pos_k).unsqueeze (1) <= torch.arange (n_doc)).int()
    
    neg_sort = neg_mask * preds  +  pos_mask * (min_preds - 1).unsqueeze (1)
    neg_sort = neg_sort.sort (dim = 1)[0]
    neg_k = (hardest_fraction * neg_mask.sum (dim = 1)).ceil().maximum (torch.ones (1))
    neg_hard_mask = ((n_doc - neg_k).unsqueeze (1) <= torch.arange (n_doc)).int()
    
    margin_mask = neg_hard_mask.unsqueeze (1) * pos_hard_mask.unsqueeze (2)
    margins = ((neg_sort.unsqueeze (1) - pos_sort.unsqueeze (2)) + margin).clamp (min = 0)
    
    losses = (margin_mask * margins).sum (dim = (1, 2)) / margin_mask.sum (dim = (1, 2))
    loss = losses.mean()
    return loss


# code with loop

def hardest_triplet_margin_loss(preds, targets, margin, hardest_fraction):
    """
        preds is a vector of N scores.
        targets is a vector of N floats being the 
	'positiveness' of the corresponding element in preds.
    """
    targets = targets >= 0.5
    positives = preds[targets]
    negatives = preds[~targets]

    hardest_positives = top_fraction(positives, hardest_fraction, False)\
            .unsqueeze(1)
    hardest_negatives = top_fraction(negatives, hardest_fraction, True)\
            .unsqueeze(0).expand(hardest_positives.numel(), -1) 

    losses = (margin + hardest_negatives - hardest_positives).clamp(min=0)
    loss = losses.mean()
    return loss

def top_fraction(x, fraction, largest=True):
    k = max(1, math.ceil(fraction * x.numel()))
    out = x.topk(k, largest=largest)[0]
    return out

num_anchors = 100
num_documents = 200

margin = 0.2
hardest_fraction = 0.5

# Matrix of similarity scores of anchors with documents (predicted by a NN)
sim = torch.rand(num_anchors, num_documents)

# Matrix of 'positiveness' of documents for each query. 
# (>=0.5 means it's a positive) (given as input in the data)
targets = torch.rand(num_anchors, num_documents)

anchor_losses = torch.zeros(num_anchors)
for i in range(num_anchors):
    anchor_losses[i] = hardest_triplet_margin_loss(
        sim[i], targets[i], margin, hardest_fraction) 

loss = anchor_losses.mean()

###

print ('loss =', loss)

lossB = hardest_triplet_margin_lossB (sim, targets, margin, hardest_fraction)
print ('lossB =', lossB)

print ('torch.allclose (lossB, loss) =', torch.allclose (lossB, loss))
print ('(lossB == loss) =', (lossB == loss))

1.7.1
loss = tensor(0.6941)
lossB = tensor(0.6941)
torch.allclose (lossB, loss) = True
(lossB == loss) = tensor(True)

Best.
K. Frank