使用 Numpy 进行外减法

Question

我只是想做：C_i=\Sum_k (A_i -B_k)^2 我看到使用简单的for loop 进行此计算比使用numpy.subtract.outer 更快！无论如何，我觉得numpy.einsum 会是最快的。我无法理解numpy.einsum。谁能帮帮我？此外，如果有人解释如何用numpy.einsum 编写由向量/矩阵组成的一般求和表达式，那就太好了？

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

带循环的 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