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

Question

我有两个numpy数组：

数组1：500,000行x100列

数组2：160,000行x100列

我想要找到每一行在数组1中与数组2之间的最大余弦相似度。换句话说，我计算数组1中第一行与数组2中所有行之间的余弦相似度，并找到最大的余弦相似度，然后计算数组1中第二行与数组2中所有行之间的余弦相似度，并找到最大的余弦相似度；对于数组1的其余部分也是如此。

我目前使用sklearn的cosine_similarity()函数并执行以下操作，但速度非常慢。我想知道是否有更快的方法来完成我想做的事情，而不涉及多进程/多线程。此外，我拥有的数组不是稀疏的。

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)))

- Alex

---

除非我误解了问题，否则你知道这总是需要对行进行80000000000次操作吗？ - Denziloe

是的...这就是为什么它很慢。任务的性质如下：Array2是160k个文档的数字表示。Array1是500k个文档的数字表示。我想找出每个500k个文档中与哪个160k个文档最相似，因此使用余弦相似度l。 - Alex

好的。我的意思只是无论如何优化，这都需要很长时间来完成。问题可能不在于cosine_similarity。 - Denziloe

这是一个有趣的问题，我会尝试解决它。 - Denziloe

说实话，考虑到你所选择的工具，你的做法似乎是个不错的方法。你正在对其中一行与整个第二个数组进行向量计算...这是一个很好的方法。也许你可以考虑一下这篇文章？https://stackoverflow.com/questions/47625437/parallel-cosine-similarity-of-two-large-files-with-each-other - kevinkayaks

@Alex 你了解向量化的好处吗？我正在测试的解决方案只涉及将A切成足够小的块（但不是行）。 - Denziloe

Answer 1

在Python中进行迭代可能会相当缓慢。尽可能使用numpy数组上的“向量化”和操作是最好的选择，这将工作传递给numpy的低级实现，速度很快。

cosine_similarity已经向量化了。因此，理想的解决方案只涉及cosine_similarity(A,B)，其中A和B是您的第一个和第二个数组。不幸的是，该矩阵为500,000乘以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)

这个操作每处理1000行A数据大约需要2秒钟的时间。因此，对于你的数据来说，需要大约1000秒的时间。

- Denziloe

---

非常感谢。这帮助我减少了在大数组上计算余弦相似度时所占用的内存空间。 - TYL

Answer 2

我只是添加了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)

- eroot163pi