python - 在 numba 中实现的 tensordot 算法比 numpy 慢得多

Question

我正在尝试扩展 numpy“tensordot”这样的东西:K_ijklm = A_ki * B_jml 可以这样写得更清晰:K = mytensordot(A,B,[2,0],[1,4,3])

据我了解，numpy 的 tensordot(带有可选参数 0)将能够执行如下操作:K_kijml = A_ki * B_jml，即保持索引的顺序。因此，我将不得不执行多次 np.swapaxes() 以获得矩阵“K_ijklm”，在复杂的情况下，这可能是一个容易出错的来源(可能很难调试)。

问题是我的实现速度很慢(比 tensordot 慢 10 倍[编辑:实际上比那个慢很多])，即使在使用 numba 时也是如此。我想知道是否有人对可以做些什么来提高我的算法的性能有一些见解。

MWE

import numpy as np
import numba as nb
import itertools
import timeit

@nb.jit()
def myproduct(dimN):
    N=np.prod(dimN)
    L=len(dimN)
    Product=np.zeros((N,L),dtype=np.int32)
    rn=0
    for n in range(1,N):
        for l in range(L):
            if l==0:
                rn=1
            v=Product[n-1,L-1-l]+rn
            rn = 0
            if v == dimN[L-1-l]:
                v = 0
                rn = 1
            Product[n,L-1-l]=v
    return Product

@nb.jit()
def mytensordot(A,B,iA,iB):
    iA,iB = np.array(iA,dtype=np.int32),np.array(iB,dtype=np.int32)
    dimA,dimB = A.shape,B.shape
    NdimA,NdimB=len(dimA),len(dimB)

    if len(iA) != NdimA: raise ValueError("iA must be same size as dim A")
    if len(iB) != NdimB: raise ValueError("iB must be same size as dim B")

    NdimN = NdimA + NdimB
    dimN=np.zeros(NdimN,dtype=np.int32)
    dimN[iA]=dimA
    dimN[iB]=dimB
    Out=np.zeros(dimN)
    indexes = myproduct(dimN)

    for nidxs in indexes:
        idxA = tuple(nidxs[iA])
        idxB = tuple(nidxs[iB])
        v=A[(idxA)]*B[(idxB)]
        Out[tuple(nidxs)]=v
    return Out



A=np.random.random((4,5,3))
B=np.random.random((6,4))

def runmytdot():
    return mytensordot(A,B,[0,2,3],[1,4])
def runtensdot():
    return np.tensordot(A,B,0).swapaxes(1,3).swapaxes(2,3)


print(np.all(runmytdot()==runtensdot()))
print(timeit.timeit(runmytdot,number=100))
print(timeit.timeit(runtensdot,number=100))

结果:

True
1.4962144780438393
0.003484356915578246

Answer 1

你遇到了a known issue . numpy.zeros requires a tuple创建多维数组时。如果你传递元组以外的东西，它有时会起作用，但这只是因为 numpy首先将对象转换为元组很聪明。

问题是numba目前不支持 conversion of arbitrary iterables into tuples .因此，当您尝试在 nopython=True 中编译它时，这一行失败了模式。 (其他几个也失败了，但这是第一次。)

Out=np.zeros(dimN)

理论上你可以调用np.prod(dimN) ，创建一个由零组成的平面数组，并对其进行整形，但随后您遇到了同样的问题:reshape numpy的方法|数组需要一个元组!

这对于 numba 来说是一个相当棘手的问题--我以前没有遇到过。我真的怀疑我找到的解决方案是否正确，但它是一个可行的解决方案，允许我们在 nopython=True 中编译一个版本模式。

核心思想是通过直接实现一个跟在strides之后的索引器来避免使用元组进行索引。数组的:

@nb.jit(nopython=True)
def index_arr(a, ix_arr):
    strides = np.array(a.strides) / a.itemsize
    ix = int((ix_arr * strides).sum())
    return a.ravel()[ix]

@nb.jit(nopython=True)
def index_set_arr(a, ix_arr, val):
    strides = np.array(a.strides) / a.itemsize
    ix = int((ix_arr * strides).sum())
    a.ravel()[ix] = val

这允许我们在不需要元组的情况下获取和设置值。

我们也可以避免使用reshape通过将输出缓冲区传递给 jitted 函数，并将该函数包装在辅助函数中:

@nb.jit()  # We can't use nopython mode here...
def mytensordot(A, B, iA, iB):
    iA, iB = np.array(iA, dtype=np.int32), np.array(iB, dtype=np.int32)
    dimA, dimB = A.shape, B.shape
    NdimA, NdimB = len(dimA), len(dimB)

    if len(iA) != NdimA:
        raise ValueError("iA must be same size as dim A")
    if len(iB) != NdimB:
        raise ValueError("iB must be same size as dim B")

    NdimN = NdimA + NdimB
    dimN = np.zeros(NdimN, dtype=np.int32)
    dimN[iA] = dimA
    dimN[iB] = dimB
    Out = np.zeros(dimN)
    return mytensordot_jit(A, B, iA, iB, dimN, Out)

由于助手不包含循环，因此它会增加一些开销，但开销非常小。这是最终的 jitted 函数:

@nb.jit(nopython=True)
def mytensordot_jit(A, B, iA, iB, dimN, Out):
    for i in range(np.prod(dimN)):
        nidxs = int_to_idx(i, dimN)
        a = index_arr(A, nidxs[iA])
        b = index_arr(B, nidxs[iB])
        index_set_arr(Out, nidxs, a * b)
    return Out

不幸的是，这并没有像我们希望的那样产生加速。在较小的阵列上，它比 tensordot 慢大约 5 倍;在更大的阵列上，它仍然慢 50 倍。 (但至少它没有慢 1000 倍!)回想起来这并不奇怪，因为 dot和 tensordot都在后台使用 BLAS，如@hpaulj reminds us .

完成这段代码后，我看到了 einsum解决了你真正的问题——很好!

但是您最初的问题所指向的潜在问题——在 jitted 代码中不可能使用任意长度的元组进行索引——仍然令人沮丧。所以希望这对其他人有用!