calculate similarity for 3-dim input

[dangvansam]

i have inputs are: anchors: torch.Size([8, 20, 128]) positives: torch.Size([8, 20, 128]) negatives: torch.Size([100, 8, 20, 128])

8-batch size, 20-num pairs, 128-embedding dim, 100-num negative samples (shuffled from positive samples)

show, can i calculate similarity for 3d inputs with your code? thanks

[RElbers]

Hi. This kind of input won't work, because the shapes don't match what the function expects. The shapes should be: anchors : (N, D), positives : (N, D), negatives: (N, M, D). If you combine the dimensions with 8 and 20, then it should work with InfoNCE(negative_mode='paired'), but I'm not sure that is what you need. Let me know if it works for you.

    anchors = torch.randn(8, 20, 128)
    positives = torch.randn(8, 20, 128)
    negatives = torch.randn(100, 8, 20, 128)

    loss = InfoNCE(negative_mode='paired')
    anchors = anchors.reshape(-1, 128)
    positives = positives.reshape(-1, 128)
    negatives = negatives.reshape(100, -1, 128).transpose(0,1)
    output = loss(anchors, positives, negatives)