How to vectorize this computation of a mask for triplet loss function

user_9781778 0
Viewed 41

Suppose I have a list of numbers lst of length N, along with two numbers epsilon and tau. I want to find the (N,N,N) mask matrix mask such that mask[i][j][k]=1 if and only if

abs(lst[i] - lst[j]) <= epsilon and abs(lst[i] - lst[k]) >= tau

This is what I tried:

d_mat = torch.cdist(lst.unsqueeze(0), lst.unsqueeze(0))
within_eps = torch.where(dmat  <= eps, 1, 0)
over_tau = torch.where(dmat >= tau, 1, 0)

mask = torch.zeros((N,N,N))
for i in range(N):
    for j in range(N):
        for k in range(N):
            if within_eps[i][j] == 1 and over_tau[i][k] == 1:
                mask[i][j][k] = 1
            else:
                mask[i][j][k] = 0

So basically I did it naively. Can you show me, with steps, how you come up with a vectorization for this?

pytorch vectorization python-3.x
1 Answers

You successfully created the 2d dmat of pairwise distances. Now you can use torch.logical_and for creating the mask:

mask = torch.logical_and(dmat[..., None] <= eps, dmat[:, None, :] >= tau)

If you want to be explicit about the distance computation (and less efficient) you can:

mask = torch.logical_and(torch.abs(lst[:, None, None] - lst[None, :, None]) <= eps,
                         torch.abs(lst[:, None, None] - lst[None, None, :]) >= tau)
user_17144100
Related