使用 numba 计算向量和矩阵行之间的余弦相似度
Using numba for cosine similarity between a vector and rows in a matix
发现 gist 使用 numba 快速计算余弦相似度。
import numba
@numba.jit(target='cpu', nopython=True)
def fast_cosine(u, v):
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for i in range(m):
if (np.isnan(u[i])) or (np.isnan(v[i])):
continue
udotv += u[i] * v[i]
u_norm += u[i] * u[i]
v_norm += v[i] * v[i]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
return ratio
结果看起来很有希望(500ns 对比只有 200us,我的机器没有 jit decorator)。
我想使用 numba 来 并行化 向量 u
和候选矩阵 M
之间的计算——即每一行的余弦。
示例:
def fast_cosine_matrix(u, M):
"""
Return array of cosine similarity between u and rows in M
>>> import numpy as np
>>> u = np.random.rand(100)
>>> M = np.random.rand(10, 100)
>>> fast_cosine_matrix(u, M)
"""
一种方法是用第二个输入重写一个矩阵。但是如果我尝试遍历矩阵的行,我会得到 NotImplementedError
。打算尝试只使用切片。
我考虑过使用 vectorize
,但我无法使用它。
稍微重写一下解决方案:
import numpy as np
import numba
@numba.jit(target='cpu', nopython=True, parallel=True)
def fast_cosine_matrix(u, M):
scores = np.zeros(M.shape[0])
for i in numba.prange(M.shape[0]):
v = M[i]
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for j in range(m):
if (np.isnan(u[j])) or (np.isnan(v[j])):
continue
udotv += u[j] * v[j]
u_norm += u[j] * u[j]
v_norm += v[j] * v[j]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
scores[i] = ratio
return scores
u = np.random.rand(100)
M = np.random.rand(100000, 100)
fast_cosine_matrix(u, M)
备选方案:使用 numba 创建通用 UFunc
@numba.guvectorize(["void(float64[:], float64[:], float64[:])"], "(n),(n)->()", target='parallel')
def fast_cosine_gufunc(u, v, result):
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for i in range(m):
if (np.isnan(u[i])) or (np.isnan(v[i])):
continue
udotv += u[i] * v[i]
u_norm += u[i] * u[i]
v_norm += v[i] * v[i]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
result[:] = ratio
u = np.random.rand(100)
M = np.random.rand(100000, 100)
fast_cosine_gufunc(u, M[0,:])
fast_cosine_gufunc(u, M)
发现 gist 使用 numba 快速计算余弦相似度。
import numba
@numba.jit(target='cpu', nopython=True)
def fast_cosine(u, v):
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for i in range(m):
if (np.isnan(u[i])) or (np.isnan(v[i])):
continue
udotv += u[i] * v[i]
u_norm += u[i] * u[i]
v_norm += v[i] * v[i]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
return ratio
结果看起来很有希望(500ns 对比只有 200us,我的机器没有 jit decorator)。
我想使用 numba 来 并行化 向量 u
和候选矩阵 M
之间的计算——即每一行的余弦。
示例:
def fast_cosine_matrix(u, M):
"""
Return array of cosine similarity between u and rows in M
>>> import numpy as np
>>> u = np.random.rand(100)
>>> M = np.random.rand(10, 100)
>>> fast_cosine_matrix(u, M)
"""
一种方法是用第二个输入重写一个矩阵。但是如果我尝试遍历矩阵的行,我会得到 NotImplementedError
。打算尝试只使用切片。
我考虑过使用 vectorize
,但我无法使用它。
稍微重写一下解决方案:
import numpy as np
import numba
@numba.jit(target='cpu', nopython=True, parallel=True)
def fast_cosine_matrix(u, M):
scores = np.zeros(M.shape[0])
for i in numba.prange(M.shape[0]):
v = M[i]
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for j in range(m):
if (np.isnan(u[j])) or (np.isnan(v[j])):
continue
udotv += u[j] * v[j]
u_norm += u[j] * u[j]
v_norm += v[j] * v[j]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
scores[i] = ratio
return scores
u = np.random.rand(100)
M = np.random.rand(100000, 100)
fast_cosine_matrix(u, M)
备选方案:使用 numba 创建通用 UFunc
@numba.guvectorize(["void(float64[:], float64[:], float64[:])"], "(n),(n)->()", target='parallel')
def fast_cosine_gufunc(u, v, result):
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for i in range(m):
if (np.isnan(u[i])) or (np.isnan(v[i])):
continue
udotv += u[i] * v[i]
u_norm += u[i] * u[i]
v_norm += v[i] * v[i]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
result[:] = ratio
u = np.random.rand(100)
M = np.random.rand(100000, 100)
fast_cosine_gufunc(u, M[0,:])
fast_cosine_gufunc(u, M)