为什么余弦相似度的矩阵计算比向量计算循环慢得多？

Question

我正在尝试优化余弦相似度计算以尽可能高效。首先，我计算向量，将所述向量存储在矩阵中，最后调用每一行来执行计算。在尝试优化时，我认为对每个向量的整个矩阵执行计算可能会更有效（删除内部 for 循环）。如果没有 jit，它的效果很好，可以减少处理所需的时间。尽管跨每个向量的 jit 仍然更快。当我将 jit 应用于矩阵计算时，处理时间激增，使其效率最低。

有谁知道为什么会发生这种情况？

模块

import numpy as np
from numba import jit
from timeit import default_timer as timer

计算代码

# @jit(nopython=True)
def cosine_sim_matrix(matrix, vector):
    dot_product = np.dot(matrix, vector)
    magnitude_a = np.sqrt((matrix ** 2).sum(axis=1))
    magnitude_b = np.sqrt((vector ** 2).sum())
    similarity_score = dot_product / (magnitude_a * magnitude_b)
    similarity_score[similarity_score > 1] = 1
    similarity_score[similarity_score < -1] = -1
    return similarity_score

# @jit(nopython=True)
def cosine_sim_vector(vector_1, vector_2):
    dot_product = np.dot(vector_1, vector_2)
    magnitude_a = np.sqrt((vector_1** 2).sum())
    magnitude_b = np.sqrt((vector_2 ** 2).sum())
    similarity_score = dot_product / (magnitude_a * magnitude_b)
    if similarity_score > 1:
        similarity_score = 1.0
    elif similarity_score < -1:
        similarity_score = -1.0 
    return similarity_score

测试

mymatrix = np.random.rand(10000,1000)

# timing matrix calculations
start = timer()
for idx in range(mymatrix.shape[0]):
    vec = mymatrix[idx,:]
    cosine_sim_matrix(mymatrix, vec)
end = timer()
print(f"process time, matrix={end-start}")

# timing vector calculations
start = timer()
for idx in range(mymatrix.shape[0]):
    vec = mymatrix[idx,:]
    for idx2 in range(mymatrix.shape[0]):
        vec2 = mymatrix[idx2,:]
        cosine_sim_vector(vec, vec2)
end = timer()
print(f"process time, vector={end-start}")

结果

# no jit
process time, matrix=348.5313815
process time, vector=839.2151422999999

# with jit
process time, matrix=883.4681065999998
process time, vector=275.0485957000001

我认为可能与

matrix ** 2

有关，并选择更改为

np.linalg.norm

，但当前的numba版本不支持轴计算。

编辑： 我发现了这篇关于改进 numba 矩阵余弦计算的文章。这里的代码非常高效，选择多个 for 循环来生成点积和幅度。

这是使用此代码的速度

process time, no parallel, matrix=109.96504489999916
process time, w/ parallel, matrix=21.85965199999964 

我的问题与有问题的代码有更多关系。

Answer 1

我不确定为什么从

cosine_sim_vector

更改为

cosine_sim_matrix

会导致 numba 速度如此之大。 尽管您应该注意，您正在计时编译时间，并且通常您还会在

njit

装饰函数中包含调用函数的 for 循环，因为它们也可以被编译。

您可以完全用 numpy 编写（除了需要重写的范数计算之外，与 numba 兼容），这比给定矩阵大小的代码要快得多：

def cosine_similarity(matrix):
    matrix_norm = np.linalg.norm(matrix, axis=1, keepdims=True)
    ratio = matrix / matrix_norm
    out = ratio.dot(ratio.T)
    out[out > 1.] = 1.
    out[out < -1.] = -1.
    return out

%timeit cosine_similarity(mymatrix)

结果：

644 ms ± 17.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)