使用多核的 Numpy np.einsum 数组乘法

Question

我已经用 MKL 编译了 numpy 1.6.2 和 scipy，希望有更好的性能。 目前我有一个严重依赖 np.einsum() 的代码，有人告诉我 einsum 不适用于 MKL，因为几乎没有矢量化。 =( 所以我想用 np.dot() 和切片重新编写我的一些代码，只是为了能够加快多核速度。 我真的很喜欢 np.einsum() 的简单性并且可读性很好。 无论如何，例如，我有一个形式的多维矩阵乘法：

np.einsum('mi,mnijqk->njqk',A,B)

那么我如何在 np.dot() 高效的 MKL 操作中转换这样的东西，或其他 3,4 和 5 维数组乘法？

我将发布更多信息： 我正在计算这个方程：

为此，我正在使用代码：

np.einsum('mn,mni,nij,nik,mi->njk',a,np.exp(b[:,:,np.newaxis]*U[np.newaxis,:,:]),P,P,X)

这不是那么快，用 cython 编码的同样的东西要快 5 倍：

    #STACKOVERFLOW QUESTION:
from __future__ import division
import numpy as np
cimport numpy as np
cimport cython

cdef extern from "math.h":
    double exp(double x)


DTYPE = np.float

ctypedef np.float_t DTYPE_t
@cython.boundscheck(False) # turn of bounds-checking for entire function
def cython_DX_h(np.ndarray[DTYPE_t, ndim=3] P, np.ndarray[DTYPE_t, ndim=1] a, np.ndarray[DTYPE_t, ndim=1] b, np.ndarray[DTYPE_t, ndim=2] U,  np.ndarray[DTYPE_t, ndim=2] X, int I, int M):
    assert P.dtype == DTYPE and a.dtype == DTYPE and b.dtype == DTYPE and U.dtype == DTYPE and X.dtype == DTYPE

cdef np.ndarray[DTYPE_t,ndim=3] DX_h=np.zeros((N,I,I),dtype=DTYPE)
cdef unsigned int j,n,k,m,i
for n in range(N):
    for j in range(I):
        for k in range(I):
            aux=0
            for m in range(N):
                for i in range(I):
                    aux+=a[m,n]*exp(b[m,n]*U[n,i])*P[n,i,j]*P[n,i,k]*X[m,i]
            DX_h[n,j,k]=aux
return DX_h

有没有办法用 cython 的性能在纯 python 中做到这一点？ （我还没有弄清楚如何张量这个方程） 无法在这个 cython 代码中执行 prange，出现很多 gil 和 nogil 错误。

Answer 1

或者，您可以使用numpy.tensordot()：

np.tensordot(A, B, axes=[[0, 1], [0, 2]])

这也将使用多个内核，例如numpy.dot()。