使用 numpy.tensordot 替换嵌套循环

Question

我有一段代码，但我想提高性能。我的代码是：

lis = []
for i in range(6):
    for j in range(6):
        for k in range(6):
            for l in range(6):
                lis[i][j] += matrix1[k][l] * (2 * matrix2[i][j][k][l] - matrix2[i][k][j][l])  
print(lis)

matrix2 是一个 4 维 np-array，matrix1 是一个二维数组。

我想通过使用 np.tensordot(matrix1, matrix2) 来加速这段代码，但后来我迷路了。

Answer 1

你可以只使用 jit 编译器

\n\n

你的解决方案一点也不差。我唯一改变的是索引和变量循环范围。\n如果你有 numpy 数组和过多的循环，你可以使用编译器（Numba），这是一件非常简单的事情。

\n\n

import numba as nb\nimport numpy as np\n#The function is compiled only at the first call (with using same datatypes)\n@nb.njit(cache=True) #set cache to false if copying the function to a command window\ndef almost_your_solution(matrix1,matrix2):\n  lis = np.zeros(matrix1.shape,np.float64)\n  for i in range(matrix2.shape[0]):\n      for j in range(matrix2.shape[1]):\n          for k in range(matrix2.shape[2]):\n              for l in range(matrix2.shape[3]):\n                  lis[i,j] += matrix1[k,l] * (2 * matrix2[i,j,k,l] - matrix2[i,k,j,l])\n\n  return lis\n

\n\n

关于代码简单性，我更喜欢 hpaulj 的 einsum 解决方案，而不是上面显示的解决方案。在我看来，tensordot 解决方案并不那么容易理解。但这是品味问题。

\n\n

性能比较

\n\n

我用于比较的 hpaulj 函数：

\n\n

def hpaulj_1(matrix1,matrix2):\n  matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3)\n  return np.einsum(\'kl,ijkl->ij\', matrix1, matrix3)\n\ndef hpaulj_2(matrix1,matrix2):\n  matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3)\n  (matrix1*matrix3).sum(axis=(2,3))\n  return np.tensordot(matrix1, matrix3, [[0,1],[2,3]])\n

\n\n

非常短的数组给出：

\n\n

matrix1=np.random.rand(6,6)\nmatrix2=np.random.rand(6,6,6,6)\n\nOriginal solution:    2.6 ms\nCompiled solution:    2.1 \xc2\xb5s\nEinsum solution:      8.3 \xc2\xb5s\nTensordot solution:   36.7 \xc2\xb5s\n

\n\n

更大的数组给出：

\n\n

matrix1=np.random.rand(60,60)\nmatrix2=np.random.rand(60,60,60,60)\n\nOriginal solution:    13,3 s\nCompiled solution:    18.2 ms\nEinsum solution:      115  ms\nTensordot solution:   180  ms\n

\n\n

结论

\n\n

编译将计算速度提高了大约 3 个数量级，并且远远优于所有其他解决方案。

\n