Search code examples
pythonpytorchtorch

Is this the right way to compute cosine similarity in PyTroch?

cos = torch.nn.CosineSimilarity(dim=-1, eps=1e-6)
c = torch.FloatTensor([1, 2, 4])
b = torch.FloatTensor([1, 2, 3])
simi = cos(b,c)
tensor(0.9915)

I am using dim=-1 in this funciton, does that mean it is a one dimension float list? Is this correct?

Solution

  • Like with most indexing in python, -1 refers to last dimension (-2 would be second-to-last, etc...). Using dim=-1 when initializing cosine similarity means that cosine similarity will be computed along the last dimension of the inputs.

    For example, if b and c were 3-dimensional tensors with size [X,Y,Z], then the result would be a 2-dimensional tensor of size [X,Y]. In your case, since the input tensors only have one dimension (size [3]), you end up getting a result tensor of size [], i.e. a scalar.