在 numba 中实现的 tensordot 算法比 numpy 慢得多
Algorithm for tensordot implemented in numba is much slower than numpy's
我正在尝试扩展 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
创建多维数组时，您将 运行 转换为 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
这使我们无需元组即可获取和设置值。
我们还可以通过将输出缓冲区传递到 jitted 函数并将该函数包装在辅助函数中来避免使用 reshape：
@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 .
完成这段代码后，我看到 einsum 已经解决了您的实际问题 -- 太好了！
但是您最初的问题所指向的根本问题——在 jitted 代码中不可能使用任意长度的元组进行索引——仍然令人沮丧。所以希望这对其他人有用！
tensordot 标量轴值可能是模糊的。我在
中进行了探索
我推导出 np.tensordot(A, B, axes=0) 等价于 axes=[[], []]。
In [757]: A=np.random.random((4,5,3))
...: B=np.random.random((6,4))
In [758]: np.tensordot(A,B,0).shape
Out[758]: (4, 5, 3, 6, 4)
In [759]: np.tensordot(A,B,[[],[]]).shape
Out[759]: (4, 5, 3, 6, 4)
这又相当于调用 dot 使用新的大小为 1 的积和维度：
In [762]: np.dot(A[...,None],B[...,None,:]).shape
Out[762]: (4, 5, 3, 6, 4)
(4,5,3,1) * (6,1,4) # the 1 is the last of A and 2nd to the last of B
dot 很快，使用 BLAS（或等效）代码。交换轴和重塑也比较快。
einsum 给了我们很多对轴的控制
复制上述产品：
In [768]: np.einsum('jml,ki->jmlki',A,B).shape
Out[768]: (4, 5, 3, 6, 4)
并交换：
In [769]: np.einsum('jml,ki->ijklm',A,B).shape
Out[769]: (4, 4, 6, 3, 5)
一个小问题——双重交换可以写成一个转置：
.swapaxes(1,3).swapaxes(2,3)
.transpose(0,3,1,2,4)
我正在尝试扩展 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
创建多维数组时，您将 运行 转换为 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
这使我们无需元组即可获取和设置值。
我们还可以通过将输出缓冲区传递到 jitted 函数并将该函数包装在辅助函数中来避免使用 reshape：
@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
完成这段代码后，我看到 einsum 已经解决了您的实际问题 -- 太好了！
但是您最初的问题所指向的根本问题——在 jitted 代码中不可能使用任意长度的元组进行索引——仍然令人沮丧。所以希望这对其他人有用！
tensordot 标量轴值可能是模糊的。我在
我推导出 np.tensordot(A, B, axes=0) 等价于 axes=[[], []]。
In [757]: A=np.random.random((4,5,3))
...: B=np.random.random((6,4))
In [758]: np.tensordot(A,B,0).shape
Out[758]: (4, 5, 3, 6, 4)
In [759]: np.tensordot(A,B,[[],[]]).shape
Out[759]: (4, 5, 3, 6, 4)
这又相当于调用 dot 使用新的大小为 1 的积和维度：
In [762]: np.dot(A[...,None],B[...,None,:]).shape
Out[762]: (4, 5, 3, 6, 4)
(4,5,3,1) * (6,1,4) # the 1 is the last of A and 2nd to the last of B
dot 很快，使用 BLAS（或等效）代码。交换轴和重塑也比较快。
einsum 给了我们很多对轴的控制
复制上述产品：
In [768]: np.einsum('jml,ki->jmlki',A,B).shape
Out[768]: (4, 5, 3, 6, 4)
并交换：
In [769]: np.einsum('jml,ki->ijklm',A,B).shape
Out[769]: (4, 4, 6, 3, 5)
一个小问题——双重交换可以写成一个转置：
.swapaxes(1,3).swapaxes(2,3)
.transpose(0,3,1,2,4)