numpy 我想提高余弦相似度计算的效率,使其更快

wrrgggsh  于 2023-06-29  发布在  其他
关注(0)|答案(1)|浏览(135)

我有一个大小为(96341,1000)的numpy数组。我想求这个数组的余弦相似度。我正在使用的机器是8 vCPU 32 GB。这是我的初始代码。而且我希望这个函数运行得更快,可以控制/限制内存的使用量,并保持相同的结果。

vec.shape
(96341,1000)

def csm(A):
    
    num = np.dot(A, A.T)  
    p1 = np.sqrt(np.sum(A**2, axis=1))[:, cp.newaxis]  
    p2 = np.sqrt(np.sum(A**2, axis=1))[cp.newaxis, :]  

    result = num / (p1 * p2)  
    
    return result

cos = csm(vec)
0x6upsns

0x6upsns1#

正如@mpw2所说,你的分子是96341×96341。但你的分母也是一样。这意味着在某个时候,你的内存中会有3× 74 GB的数组(一个保存分子num,一个临时保存分母p1*p2,一个保存结果)。
对于分母,这是很容易解决的,我看不出有任何理由不这样做

def csm(A):
    
    num = np.dot(A, A.T)  
    p1 = np.sqrt(np.sum(A**2, axis=1))[:, cp.newaxis]  
    p2 = np.sqrt(np.sum(A**2, axis=1))[cp.newaxis, :]  

    # Only change. Let the broadcasting (which does not allocate 
    # memory: p1 and p2 are only "implicit" 96341x96341 arrays)
    # does the job.
    # Note that the parenthesis are useless here, since / are evaluated
    # from left to right anyway. I just want it to be explicit  
    result = (num / p1)/p2
    
    return result

还要注意,p1p2是相同的数组,只是视图发生了变化。所以不需要计算两次

def csm(A):
    
    num = np.dot(A, A.T)  
    px = np.sqrt(np.sum(A**2, axis=1))
    result = num / px[:,None] / px[None,:]
    
    return result

最后,如果我们想避免那些74 GB(仍然是2x 74 GB,代码为:我们仍然有num和result在同一时刻都在内存中),我们可以按块计算

def csm2(A,B):
    num = np.dot(A, B.T)  
    # this time they are not the same
    pa = np.sqrt(np.sum(A**2, axis=1))[:,None]
    pb = np.sqrt(np.sum(B**2, axis=1))[None, :]
    return num / pa / pb

然后你可以按块计算

res=np.empty((len(vec), len(vec)))
N=1000 # Block size (it doesn't need to be a divisor of len(vec). Batch at end of rows or columns will be smaller that's all
for i in range(0, len(vec), N):
    for j in range(0, len(vec), N):
        res[i:i+N, j:j+N] = csm2(vec[i:i+N], vec[j:j+N])
        res[j:j+N, i:i+N] = res[i:i+N, j:j+N].T

我们不仅避免了大部分的冗余计算,因为对称性已经提到。但是我们在内存中从来没有大于NxN的数组,但是对于res本身(但是如果这是你想要的结果,这是不可避免的)。
当然,我们使用for循环,几乎所有我的答案都提到for循环是numpy中的终极罪行。但是这些都是可以的,因为如果N足够大,那么迭代所花费的计算在numpy函数的计算时间之前将是可以忽略的。
最后剩下的问题是,正如我所说,结果的74 GB(至少,现在,只有一个,而不是你的解决方案中的3个,而不是我的第一个解决方案中的2个)。您可以在应用程序中避免这种情况,具体取决于您对结果的处理。尝试一次只使用结果的一部分。
例如,如果想法是找到最小值,则可以计算该矩阵的块,并为每个块计算最小值,然后计算min的min。这是一个经典的map/reduce解决方案
所以,与其

M=csm(vec)
print(M.min())

mn=np.inf
for i in range(0, len(vec), N):
    for j in range(i, len(vec), N):
        mn=min(mn, csm2(vec[i:i+N], vec[j:j+N]).min())
print(mn)

当然,我怀疑你想用那个csm做什么就像min一样简单。但是既然你没有说,我只能推测有办法做到这一点,同时一次在一个块上计算csm。

相关问题