Python:两个大型numpy数组之间的余弦相似度

Python: Cosine similarity between two large numpy arrays

我有两个 numpy 数组:

数组 1:500,000 行 x 100 列

数组 2:160,000 行 x 100 列

我想找到数组 1 和 [=24= 中的每一行之间的最大余弦相似度 ]数组2。换句话说,我计算数组 1 中第一行与数组 2 中所有行之间的余弦相似度,并找到最大余弦相似度,然后我计算数组 1 中第二行与数组 2 中所有行之间的余弦相似度数组2,求最大余弦相似度;并对数组 1 的其余部分执行此操作。

我目前使用 sklearncosine_similarity() 功能并执行以下操作,但速度非常慢。我想知道是否有一种不涉及 multiprocessing/multithreading 的更快方法来完成我想做的事情。另外,我的数组并不稀疏。

from sklearn.metrics.pairwise import cosine_similarity as cosine

results = []
for i in range(Array1.shape[0]):
     results.append(numpy.max(cosine(Array1[None,i,:], Array2)))

Python 中的迭代可能会很慢。总是最好 "vectorise" 并尽可能在数组上使用 numpy 操作,这会将工作传递给 numpy 的低级实现,速度很快。

cosine_similarity 已经矢量化。因此,理想的解决方案只涉及 cosine_similarity(A, B),其中 A 和 B 是您的第一个和第二个数组。不幸的是,这个矩阵是 500,000 x 160,000,这在内存中太大了(它会抛出错误)。

下一个最佳解决方案是将 A(按行)拆分为大块(而不是单独的行),以便结果仍然适合内存,并迭代它们。我发现您的数据在每个块中使用 100 行适合内存;更多,它不工作。然后我们简单地使用 .max 并为每次迭代获得 100 个最大值,我们可以在最后一起收集。

不过,这种方式强烈建议我们节省更多时间。两个向量的余弦相似度的公式为u.v / |u||v|,是两者夹角的余弦值。因为我们在迭代,所以我们每次都会重新计算 B 的行的长度并将结果丢弃。一个很好的解决方法是利用余弦相似度在缩放向量(角度相同)时不会变化的事实。所以我们可以只计算一次所有行的长度,然后除以它们,使行成为单位向量。然后我们简单地计算余弦相似度 u.v,这可以通过矩阵乘法对数组进行。我对此进行了快速测试,速度大约快了 3 倍。

综合来看:

import numpy as np

# Example data
A = np.random.random([500000, 100])
B = np.random.random([160000, 100])

# There may be a proper numpy method for this function, but it won't be much faster.
def normalise(A):
    lengths = (A**2).sum(axis=1, keepdims=True)**.5
    return A/lengths

A = normalise(A)
B = normalise(B)

results = []

rows_in_slice = 100

slice_start = 0
slice_end = slice_start + rows_in_slice

while slice_end <= A.shape[0]:

    results.append(A[slice_start:slice_end].dot(B.T).max(axis=1))

    slice_start += rows_in_slice
    slice_end = slice_start + rows_in_slice

result = np.concatenate(results)

每 1,000 行 A 到 运行 大约需要 2 秒。所以你的数据应该是大约 1,000 秒。

只需添加 numba 版本即可转换为快速机器码。

我做了很多 for 循环,因为 numpy 使用广播,它将分配临时内存,我猜它已经是内存限制了。

我刚刚在 numba 中重写了余弦逻辑。您也可以通过在 njit 选项中添加 parallel=True 来并行化它。

虽然numba是否比numpy表现更好取决于问题,但是numpy并行是困难的

import numpy as np
import numba as nb

A_1 = np.random.random((500, 100))
A_2 = np.random.random((160, 100))

@nb.njit((nb.float64[:, ::100], nb.float64[:, ::100]))
def max_cos(a, b):
    norm_a = np.empty((a.shape[0],), dtype=np.float64)
    norm_b = np.empty((b.shape[0],), dtype=np.float64)

    for i in nb.prange(a.shape[0]):
        sq_norm = 0.0
        for j in range(100):
            sq_norm += a[i][j] ** 2
        norm_a[i] = sq_norm ** 0.5
    
    for i in nb.prange(b.shape[0]):
        sq_norm = 0.0
        for j in range(100):
            sq_norm += b[i][j] ** 2
        norm_b[i] = sq_norm ** 0.5
        
    max_pair = (0, 0)
    min_dot = 1e+307
    for i in nb.prange(a.shape[0]):
        max_j = 0
        min_idot = 1e+307
        for j in range(b.shape[0]):
            dot_ij = 0.0
            for k in range(100):
                dot_ij += a[i][k] * b[j][k]
            dot_ij /= norm_b[j]
            if min_idot > dot_ij:
                min_idot = dot_ij
                max_j = j
        min_idot /= norm_a[i]
        if min_dot > min_idot:
            min_dot = min_idot
            max_pair = (i, j)
    return max_pair
%%timeit
max_cos(A_1, A_2)
# 6.03 ms ± 34 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%%timeit
from sklearn.metrics.pairwise import cosine_similarity as cosine

results = []
for i in range(A_1.shape[0]):
     results.append(np.max(cosine(A_1[None,i,:], A_2)))
# 115 ms ± 2.28 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)