大矩阵上的余弦距离

Question

我有一个大小为 (100000, 100) 的嵌入矩阵。 我想计算矩阵中所有成对的余弦距离。我尝试过使用

sklearn.metrics.pairwise.cosine_distances

函数，但由于 RAM 内存达到极限而崩溃。 我也尝试像这样批量计算：

from sklearn.metrics.pairwise import cosine_distances
embeddings.astype(np.float32)
distances_matrix = []
batch_size = 1000
df_size = len(embeddings)
for i in tqdm(range(0, df_size, batch_size)):
   end = min(i + batch_size, df_size)
   batch = embeddings[i:end]
   batch_distances = cosine_distances(batch, embeddings) 
   distances_matrix.append(batch_distances)

但是在大约 11 次迭代后它也会崩溃。

关于如何解决这个问题有什么建议吗？ 谢谢。

Answer 1

假设您正在使用 Hugging Face 进行工作并且已经使用该库进行嵌入，您应该采用他们推荐的方法并使用 FAISS

如果没有，你也可以尝试：

• 使嵌入变得稀疏。在 sklearn 中，在余弦相似度计算中使用稀疏矩阵的效率要高得多。
• 尝试批量计算并使用更简单的函数，例如scipy余弦距离
• 有人为 sklearn 中的高效成对余弦相似性编写了一个包装库，也可能值得一试：effcosim