torch.cdist高效计算大矩阵相似度

问题定义

现有矩阵$A\in R^{N\times C},B\in R^{M\times C}$，需要计算矩阵$A$和$B$的相似度（欧式距离）矩阵$S\in R^{N\times M}$，$N$和$M$很大。可以使用pytorch提供的torch.cdist方法，记得使用GPU计算。

import torch

N, M, C = 20000, 50000, 128
A = torch.rand((N, C)).cuda()
B = torch.rand((M, C)).cuda()

S = torch.cdist(A, B, p=2)
print(S.shape)