使用 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 數組，而 matrix1 是一個 2d 數組。

我想通過使用 np.tensordot(matrix1, matrix2) 來加速這段代碼，但后來我迷路了。

Answer 1

您可以只使用 jit 編譯器

你的解決方案一點也不差。 我唯一改變的是索引和可變循環范圍。 如果你有 numpy 數組和過多的循環，你可以使用編譯器（ Numba ），這是一件非常簡單的事情。

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

  return lis

關於代碼的簡單性，我更喜歡 hpaulj 的 einsum 解決方案，而不是上面顯示的解決方案。 在我看來，張量點解決方案並不容易理解。 但這是一個品味問題。

比較性能

我用於比較的 hpaulj 函數：

def hpaulj_1(matrix1,matrix2):
  matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3)
  return np.einsum('kl,ijkl->ij', matrix1, matrix3)

def hpaulj_2(matrix1,matrix2):
  matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3)
  (matrix1*matrix3).sum(axis=(2,3))
  return np.tensordot(matrix1, matrix3, [[0,1],[2,3]])

非常短的數組給出：

matrix1=np.random.rand(6,6)
matrix2=np.random.rand(6,6,6,6)

Original solution:    2.6 ms
Compiled solution:    2.1 µs
Einsum solution:      8.3 µs
Tensordot solution:   36.7 µs

更大的數組給出：

matrix1=np.random.rand(60,60)
matrix2=np.random.rand(60,60,60,60)

Original solution:    13,3 s
Compiled solution:    18.2 ms
Einsum solution:      115  ms
Tensordot solution:   180  ms

結論

編譯將計算速度提高了大約 3 個數量級，並且大大優於所有其他解決方案。

Answer 2

測試設置：

In [274]: lis = np.zeros((6,6),int)
In [275]: matrix1 = np.arange(36).reshape(6,6)
In [276]: matrix2 = np.arange(36*36).reshape(6,6,6,6)
In [277]: 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] - mat
     ...: rix2[i,k,j,l])
     ...:                 
In [278]: lis
Out[278]: 
array([[-51240,  -9660,  31920,  73500, 115080, 156660],
       [ 84840, 126420, 168000, 209580, 251160, 292740],
       [220920, 262500, 304080, 345660, 387240, 428820],
       [357000, 398580, 440160, 481740, 523320, 564900],
       [493080, 534660, 576240, 617820, 659400, 700980],
       [629160, 670740, 712320, 753900, 795480, 837060]])

正確的？

我不確定 tensordot 是不是正確的工具。 至少可能不是最簡單的。 它當然無法處理matrix2的差異。

讓我們從一個明顯的替換開始：

In [279]: matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3)
In [280]: lis = np.zeros((6,6),int)
In [281]: 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] * matrix3[i,j,k,l]

測試正常 - 相同的lis 。

現在很容易用einsum表達這一點 - 只需復制索引

In [284]: np.einsum('kl,ijkl->ij', matrix1, matrix3)
Out[284]: 
array([[-51240,  -9660,  31920,  73500, 115080, 156660],
       [ 84840, 126420, 168000, 209580, 251160, 292740],
       [220920, 262500, 304080, 345660, 387240, 428820],
       [357000, 398580, 440160, 481740, 523320, 564900],
       [493080, 534660, 576240, 617820, 659400, 700980],
       [629160, 670740, 712320, 753900, 795480, 837060]])

兩個軸上的元素乘積加和也可以； 和一個等效的tensordot量點（指定要對哪些軸求和）

(matrix1*matrix3).sum(axis=(2,3))
np.tensordot(matrix1, matrix3, [[0,1],[2,3]])

編輯

也可以使用較新的np.matmul/@ ，但需要進行一些重塑

In [111]: (matrix1.ravel()[None,None,None,:]@matrix3.reshape(6,6,-1,1)).squeeze(
     ...: )
Out[111]: 
array([[-51240,  -9660,  31920,  73500, 115080, 156660],
       [ 84840, 126420, 168000, 209580, 251160, 292740],
       [220920, 262500, 304080, 345660, 387240, 428820],
       [357000, 398580, 440160, 481740, 523320, 564900],
       [493080, 534660, 576240, 617820, 659400, 700980],
       [629160, 670740, 712320, 753900, 795480, 837060]])

這將kl維度減少到 1，並在ij維度上進行“廣播”。