Python 在大矩阵中查找相似向量的最快方法

Question

我正在尝试编写一个 python 函数，它接受一个向量 (1x128)，然后在一个大型未排序矩阵 (2000x128) 中找到最相似的列。此函数在应用程序中调用了约 100000 次。当我在桌面 PC 上工作时没有问题，但它在 Raspberry Pi 上运行非常慢。这是我的功能；

def find_similar_index(a):
    d = []
    norma=np.linalg.norm(a)
    for i in range(0, 1999):
        d.append(np.abs(np.linalg.norm(a - A[:, i]))/norma)
    return np.argmin(d)

我可以改进此功能中的任何内容以加快工作速度吗？我可以使用 Raspberry Pi 的 GPU 进行这种计算吗？

Answer 1

这是使用 broadcasting 的一种方法和 np.einsum -

subs = (a[:,None] - A)
sq_dist = np.einsum('ij,ij->j',subs, subs)
min_idx = np.abs(sq_dist).argmin()

使用 (a-b)^2 = a^2 + b^2 - 2ab 公式获取 sq_dist 的另一种方法 -

sq_dist = (A**2).sum(0) + a.dot(a) - 2*a.dot(A) 

有了 np.einsum，这会提升到 -

sq_dist = np.einsum('ij,ij->j',A,A) + a.dot(a) - 2*a.dot(A)

此外，由于我们只对最接近的索引以及与 np.linalg.norm 的距离感兴趣的最终结果在给定输入数组中的实数的情况下为正，因此我们可以跳过 np.abs 并跳过 norma 的缩放。

运行时测试

方法-

def app0(a,A): # Original approach
    d = []
    for i in range(0, A.shape[1]):
        d.append(np.linalg.norm(a - A[:, i]))
    return np.argmin(d)

def app1(a,A):
    subs = (a[:,None] - A)
    sq_dist = np.einsum('ij,ij->j',subs, subs)
    return sq_dist.argmin()

def app2(a,A):
    sq_dist = (A**2).sum(0) + a.dot(a) - 2*a.dot(A) 
    return sq_dist.argmin()

def app3(a,A):
    sq_dist = np.einsum('ij,ij->j',A,A) + a.dot(a) - 2*a.dot(A)
    return sq_dist.argmin()

既然您提到向量的形状是(1x128)，并且您正在A中寻找与该向量相似的列，所以似乎每列的长度都是 128，因此我假设 A 的形状是 (128, 2000)。根据这些假设，这里是使用列出的方法的设置和时间安排 -

In [194]: A = np.random.rand(128,2000)
     ...: a = np.random.rand(128)
     ...: 

In [195]: %timeit app0(a,A)
100 loops, best of 3: 9.21 ms per loop

In [196]: %timeit app1(a,A)
1000 loops, best of 3: 330 µs per loop

In [197]: %timeit app2(a,A)
1000 loops, best of 3: 287 µs per loop

In [198]: %timeit app3(a,A)
1000 loops, best of 3: 291 µs per loop

In [200]: 9210/287.0 # Speedup number
Out[200]: 32.09059233449477