CODEWITHSUNDEEP
pytorch
Mpizos Dimitris
Mpizos Dimitris

Reputation: 4991

torch row wise cosinus similarity

Assuming I am having 2 tensors with dimension:

and I would like to use torch.cosine_similarity to produce a tensor C of dimensions [512, 1, 200] or [512, 200].

Using torch.cosine_similarity(A, B) I get an error:

'{RuntimeError}The size of tensor a (512) must match the size of tensor b (200) at non-singleton dimension 1'

I guess it can be done by the following:

desired_result = torch.stack([torch.cosine_similarity(A_row, B_row, axis=-1) for B_row, A_row in zip(B, A)])

but there should be a more optimized way. Any help/hint?

Upvotes: 0

Views: 274

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24701

No need for stacking, broadcasting will do the job for you:

# (512, 100, 1) after unsqueeze
A = torch.randn(512, 100).unsqueeze(dim=-1)
B = torch.randn(512, 200, 100)

# (512, 200)
torch.cosine_similarity(B, A.unsqueeze(1), axis=-1)

Upvotes: 2

Related Questions