我有一段代码，但是我想提高性能。我的代码是：
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）加快此代码的速度，但是后来我迷路了。
您可以只使用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解决方案。我认为，tensordot解决方案不是那么容易理解。但这是一个品味问题。
性能比较
我用于比较的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个数量级，并且比所有其他解决方案都快了许多。
