I am using toch.bmm to compute the paired-wise cosine distance between BxDxN and BxDxN. It will return a matrix size of NxN instead of a triangle vector in the matrix in the nn.CosineSimilarity. How to use nn.CosineSimilarity to get full cosine matrix as torch.bmm did? I cannot use torch.bmm because of CUDA memory error. This is my code.
input1 = torch.randn(2, 4, 4)
input2 = torch.randn(2, 4, 4)
#Using bmm
x_norm = input1 / torch.norm(input1, p=2, dim=1, keepdim=True)
y_norm = input2 / torch.norm(input2, p=2, dim=1, keepdim=True)
cosine_sim = torch.bmm(x_norm.transpose(2,1), y_norm)
print('Using bmm: \n', cosine_sim)
# Pytorch CosineSimilarity
cos = nn.CosineSimilarity(dim=1, eps=1e-6)
cosine_sim = cos(input1, input2)
print('Using nn.CosineSimilarity: \n', cosine_sim)
The output is
Using bmm:
tensor([[[-0.0230, 0.2983, 0.0487, 0.3974],
[-0.5747, 0.5513, -0.6436, -0.1389],
[-0.3876, -0.2107, 0.7093, -0.4929],
[-0.3446, -0.5347, 0.6372, -0.6423]],
[[-0.3842, -0.0349, 0.1621, 0.6400],
[ 0.6776, -0.4812, -0.3169, -0.7976],
[-0.5251, -0.1258, 0.9381, -0.2379],
[-0.1517, 0.7164, 0.8332, 0.1668]]])
Using nn:
tensor([[-0.0230, 0.5513, 0.7093, -0.6423],
[-0.3842, -0.4812, 0.9381, 0.1668]])