Generating a bool mask by labels

Hi,

I am trying to generate a mask to get the positive logits.

I get label_1 [1,2,3,4,7] and label_2 [1,2,7,2]

the goal is to get a 01 mask M, when i and j are the same class, Mij=1, others = 0

the above label_1 and label_2 could generate such mask with size[5,4]
[[1,0,0,0],
[0,1,0,1],
[0,0,0,0],
[0,0,0,0],
[0,0,1,0]]

Well, I do generate the mask by iteration, but I it is not efficient when labels get larger.

A comparison with broadcasting should work:

label_1 = torch.tensor([1,2,3,4,7])
label_2 = torch.tensor([1,2,7,2])

res = label_1.unsqueeze(1) == label_2
print(res)
# tensor([[ True, False, False, False],
#         [False,  True, False,  True],
#         [False, False, False, False],
#         [False, False, False, False],
#         [False, False,  True, False]])

print(res.long())
# tensor([[1, 0, 0, 0],
#         [0, 1, 0, 1],
#         [0, 0, 0, 0],
#         [0, 0, 0, 0],
#         [0, 0, 1, 0]])
1 Like

Really cool! Thanks a lot.