使用numpy.tensordot替换嵌套循环 [英] Use numpy.tensordot to replace a nested loop

问题描述

我有一段代码，但是我想提高性能.我的代码是:

I have a piece of code, but I want to pull up the performance. My code is:

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数组.

matrix2 is a 4-dimensional np-array, and matrix1 is a 2d-array.

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

I want to speed up this code by using np.tensordot(matrix1, matrix2), but then I'm lost.

推荐答案

您可以只使用jit编译器

您的解决方案一点也不差.我唯一更改的是索引和变量循环范围. 如果您有numpy数组和过多的循环，则可以使用编译器( Numba )，这确实很简单去做.

Your solution isn't bad at all. The only thing I have changed is the indexing and variable loop ranges. If you have numpy arrays and excessive looping you can use a compiler (Numba), which is a really simple thing to do.

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解决方案不是那么容易理解.但这是一个品味问题.

Regarding code simplicity I would prefer the einsum solution from hpaulj over the solution shown above. The tensordot solution isn't that easy to understand to my opinion. But that's a a matter of taste.

性能比较

我用于比较的hpaulj函数:

The function from hpaulj i used for comparison:

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个数量级，并且比所有其他解决方案都快了很多.

Compilation speeds up the computation by about 3 orders of magnitude and outperforms all other solutions by quite a margin.

这篇关于使用numpy.tensordot替换嵌套循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！