python - 使用 Numpy 进行外减法

Question

我只是想做:C_i=\Sum_k (A_i -B_k)^2我发现使用简单的 for 循环 计算比使用 numpy.subtract.outer 更快!无论如何，我觉得 numpy.einsum 是最快的。我不太能理解 numpy.einsum 。有人可以帮我吗？此外，如果有人解释如何用 numpy.einsum 编写由向量/矩阵组成的通用求和表达式，那就太好了？

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

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

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 无法按 block 进行计算，这会导致大量缓存未命中，从而导致性能不佳。某些 Fortran 或 C 编译器应该能够自动执行至少以下优化(逐 block 评估)。

大型数组的实现

@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