使用 Numpy 进行外减法

Question

我只想做：C_i=\Sum_k (A_i -B_k)^2 我看到这个计算使用简单的 for loop 比使用 numpy.subtract.outer 更快！反正我觉得 numpy.einsum 会是最快的。我无法理解numpy.einsum那么好。谁能帮帮我？此外，如果有人解释如何用 numpy.einsum?

编写由 vector/matrices 组成的一般求和表达式，那就太好了

我没有在网上找到这个特定问题的解决方案。抱歉，如果以某种方式重复。

带循环的 MWE 和 numpy.subtract.outer--

A)带循环

import timeit
code1="""
import numpy as np

N=10000;

a=np.random.rand(N); b=10*(np.random.rand(N)-0.5);

def A1(x,y):
    Nx=len(x)
    z=np.zeros(Nx)
    for i in np.arange(Nx):
        z[i]=np.sum((x[i]-y)*(x[i]-y))

    return z

C1=A1(a,b)"""
elapsed_time = timeit.timeit(code1, number=10)/10
print "time=", elapsed_time

B) 与 numpy.subtract.outer

import timeit
code1="""
import numpy as np

N=10000;

a=np.random.rand(N); b=10*(np.random.rand(N)-0.5);

def A2(x,y):
    C=np.subtract.outer(x,y);
    return np.sum(C*C, axis=1)

C2=A2(a,b)"""
elapsed_time = timeit.timeit(code1, number=10)/10
print "time=", elapsed_time

对于 N=10000，循环变得更快。对于 N=100，外部减法变得更快。对于 N=10^5，外减法在我的 8GB 内存桌面上面临内存问题！

Answer 1

至少使用 Numba 或 Fortran 实现

你的两个函数都很慢。 Python 循环非常低效 (A1)，分配大型临时数组也很慢（A2 和部分 A1）。

用于小型数组的 Naive Numba 实现

import numba as nb
import numpy as np

@nb.njit(parallel=True, fastmath=True)
def A_nb_p(x,y):

    z=np.empty(x.shape[0])
    for i in nb.prange(x.shape[0]):
        TMP=0.
        for j in range(y.shape[0]):
            TMP+=(x[i]-y[j])**2
        z[i]=TMP

    return z

计时

import time
N=int(1e5)
a=np.random.rand(N)
b=10*(np.random.rand(N)-0.5)

t1=time.time()
res_1=A1(a,b)
print(time.time()-t1)
#95.30195426940918 s

t1=time.time()
res_2=A_nb_p(a,b)
print(time.time()-t1)
#0.28573083877563477 s

#A2 is too slow to measure

如上所述，这是对较大数组的简单实现，因为 Numba 无法按块进行计算，这会导致大量缓存未命中，从而导致性能不佳。某些 Fortran 或 C- 编译器应该至少能够自动执行以下优化（按块计算）。

大型数组的实现

@nb.njit(parallel=True, fastmath=True)
def A_nb_p_2(x,y):
    blk_s=1024
    z=np.zeros(x.shape[0])
    num_blk_x=x.shape[0]//blk_s
    num_blk_y=y.shape[0]//blk_s

    for ii in nb.prange(num_blk_x):
        for jj in range(num_blk_y):
            for i in range(blk_s):
                TMP=z[ii*blk_s+i]
                for j in range(blk_s):
                    TMP+=(x[ii*blk_s+i]-y[jj*blk_s+j])**2
                z[ii*blk_s+i]=TMP

    for i in nb.prange(x.shape[0]):
        TMP=z[i]
        for j in range(num_blk_y*blk_s,y.shape[0]):
            TMP+=(x[i]-y[j])**2
        z[i]=TMP

    for i in nb.prange(num_blk_x*blk_s,x.shape[0]):
        TMP=z[i]
        for j in range(num_blk_y*blk_s):
            TMP+=(x[i]-y[j])**2
        z[i]=TMP

    return z

计时

N=int(2*1e6)
a=np.random.rand(N)
b=10*(np.random.rand(N)-0.5)

t1=time.time()
res_1=A_nb_p(a,b)
print(time.time()-t1)
#298.9394392967224

t1=time.time()
res_2=A_nb_p_2(a,b)
print(time.time()-t1)
#70.12

使用 Numpy 进行外减法

Outer subtraction with Numpy

python

arrays

numpy

numpy-einsum

至少使用 Numba 或 Fortran 实现