我有一个包含64000个嵌入的数组A和另一个包含12000个嵌入的数组B(每个嵌入都是1024个浮点数)。
现在我想计算数组A和数组B之间所有对的余弦相似度(笛卡尔积)。
为了执行这个操作(使用Pandas),我使用.merge(how="cross")
将数组A和数组B合并,得到768 000 000对。
现在我正在寻找计算余弦sim的最快方法。现在我使用了类似的东西使用Numpy:
def compute_cosine_sim(a, b):
return dot(a, b)/(norm(a)*norm(b))
np.vectorize(compute_cosine_sim)(df.embedding_A.to_numpy(), df.embedding_B.to_numpy())
为了使RAM保持在合理的水平,我使用panda Dataframe分块。
问题是我的方法不够快,我想知道这里是否有什么需要改变的,特别是关于我使用的numpy函数的有效性。
给一些细节,我达到了130000 iter/sec的这个函数,这是正常的吗?还有,这种操作可以很容易地在GPU上运行吗?
谢谢你的帮助
1条答案
按热度按时间f8rj6qna1#
你可以将嵌入A与嵌入B的转置矩阵相乘,得到所有对的点积,然后将输出沿着列除以A中向量的范数,沿行除以B中向量的范数: