计算大型矩阵的余弦相似度时进程终止

Question

# A: 14287 * 768 array, B: 863394 * 768 array
def cosine_similarity(A,B):
    A = torch.tensor(A).to('cpu') ; B = torch.tensor(B).to('cpu')
    num = torch.mm(A, B.T)
    p1 = torch.sqrt(torch.sum(A**2, axis=1))[:, None]
    p2 = torch.sqrt(torch.sum(B**2, axis=1))[None, :]
    return (num/(p1*p2)).T

当我对两个矩阵进行余弦相似度计算时，进程被终止。我在机器上得到以下日志：

[Apr 5 14:38] python invoked oom-killer: gfp_mask=0x100dca(GFP_HIGHUSER_MOVABLE|__GFP_ZERO), order=0, oom_score_adj=0
[  +0.000003]  oom_kill_process+0xe6/0x120
[  +0.000001] oom-kill:constraint=CONSTRAINT_NONE,nodemask=(null),cpuset=/,mems_allowed=0,global_oom,task_memcg=/user.slice,task=python,pid=8103,uid=1000
[  +0.000013] Out of memory: Killed process 8103 (python) total-vm:81517744kB, anon-rss:61879912kB, file-rss:4kB, shmem-rss:0kB, UID:1000 pgtables:154388kB oom_score_adj:0

如果我使用 np.array 而不是 torch.tensor，我仍然会收到如下 OOM 错误：

  File "cosine.py", line 11, in cosine_similarity
    num=np.dot(A,B.T)
  File "<__array_function__ internals>", line 6, in dot
MemoryError: Unable to allocate 91.9 GiB for an array with shape (14287, 863394) and data type float64

Answer 1

这就是 Linux oom (out of memory) killer，它会选择在可用内存严重不足时终止的进程。您可以通过添加 RAM 或交换 space 或重组计算以处理较小的数据块来避免它。

Answer 2

最终使用 FAISS 库，该库通过使用其高效的搜索算法和利用 GPU 显着减少了时间。貌似2天的余弦相似度计算不到一个小时就搞定了

计算大型矩阵的余弦相似度时进程终止

Process killed while computing cosine similarity for large matrices

python

process

matrix

cosine-similarity

tensor