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

Question

我有两个 numpy 数组：

数组 1：500,000 行 x 100 列

数组 2：160,000 行 x 100 列

我想找出 数组 1 中的每一行 和 数组 2 之间的最大余弦相似度。换句话说，我计算数组 1 中的第一行和数组 2 中的所有行之间的余弦相似度，并找到最大的余弦相似度，然后我计算数组 1 中的第二行和数组 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)))

Answer 1

在 Python 中迭代可能会很慢。最好是“向量化”并尽可能对数组使用 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)

运行每 1000 行 A 大约需要 2 秒。所以你的数据应该是 1000 秒左右。

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)