【问题标题】：Python, numpy, einsum multiply a stack of matricesPython，numpy，einsum乘以一堆矩阵
【发布时间】：2014-11-16 20:07:06
【问题描述】：

出于性能原因，

我很好奇是否有一种将一堆矩阵相乘的方法。我有一个 4 维数组（500、201、2、2）。它基本上是一个 500 长度的 (201,2,2) 矩阵堆栈，对于 500 个矩阵中的每一个，我想使用 einsum 将相邻矩阵相乘并得到另一个 (201,2,2) 矩阵。

我只在最后对 [2x2] 矩阵进行矩阵乘法。由于我的解释已经偏离轨道，我将只展示我现在在做什么，以及“减少”等价物以及为什么它没有帮助（因为它的计算速度相同）。最好这将是一个麻木的单线，但我不知道那是什么，或者即使它可能。

代码：

Arr = rand(500,201,2,2)

def loopMult(Arr):
    ArrMult = Arr[0]
    for i in range(1,len(Arr)):
        ArrMult = np.einsum('fij,fjk->fik', ArrMult, Arr[i])
    return ArrMult

def myeinsum(A1, A2):
    return np.einsum('fij,fjk->fik', A1, A2)

A1 = loopMult(Arr)
A2 = reduce(myeinsum, Arr)
print np.all(A1 == A2)

print shape(A1); print shape(A2)

%timeit loopMult(Arr)
%timeit reduce(myeinsum, Arr)

True
(201, 2, 2)
(201, 2, 2)
10 loops, best of 3: 34.8 ms per loop
10 loops, best of 3: 35.2 ms per loop

任何帮助将不胜感激。东西是功能性的，但是当我必须对大量参数进行迭代时，代码往往会花费很长时间，我想知道是否有办法避免通过循环进行 500 次迭代。强>

【问题讨论】：

只是为了让您知道，使用 np.cumprod(np.einsum('fij,afjk->afik', Arr[0], Arr), axis=0)[-1] 会得到相同的结果，但您的解决方案仍然更快...也许这可以给您一些见解...
我从来没有使用过它，而且我对它的确切作用的理解是模糊的，所以这甚至可能没有意义......但是你看过 Theano 吗？
如果Arr是(n,2,2)，这个计算是Arr[0].dot(Arr[1]).dot(Arr[2])...
如果A 是(5,2,2)，那么这个计算是np.einsum('ij,jk,kl,lm,mn', *A)
我认为所有这些答案都是我问题的核心。感谢您提供有关 cumprod 的提示。它很干净，但没有速度改进。至于 Theano，我认为这超出了我的范围。 hpaulj，您的回答似乎暗示没有以这种方式使用 einsum 的动态方式。如果我知道数组的第一个维度（我并不总是这样），那么我可以为 einsum 索引硬编码一个字符串，但可能会受到字母表的限制，不是吗？如果我不知道索引，我将不得不动态生成一个索引字符串，并且我再次受到字母表的限制。

标签： python arrays performance numpy multiplication

【解决方案1】：

我认为使用 numpy 无法有效地做到这一点（不过，cumprod 解决方案很优雅）。在这种情况下，我会使用f2py。这是调用我所知道的更快语言的最简单方法，并且只需要一个额外的文件。

fortran.f90：

subroutine multimul(a, b)
  implicit none
  real(8), intent(in)  :: a(:,:,:,:)
  real(8), intent(out) :: b(size(a,1),size(a,2),size(a,3))
  real(8) :: work(size(a,1),size(a,2))
  integer i, j, k, l, m
  !$omp parallel do private(work,i,j)
  do i = 1, size(b,3)
    b(:,:,i) = a(:,:,i,size(a,4)) 
    do j = size(a,4)-1, 1, -1
      work = matmul(b(:,:,i),a(:,:,i,j))
      b(:,:,i) = work
    end do
  end do
end subroutine

使用f2py -c -m fortran fortran.f90（或F90FLAGS="-fopenmp" f2py -c -m fortran fortran.f90 -lgomp 以启用OpenMP 加速）进行编译。然后你会在你的脚本中使用它

import numpy as np, fmuls
Arr = np.random.standard_normal([500,201,2,2])
def loopMult(Arr):
  ArrMult = Arr[0]
  for i in range(1,len(Arr)):
    ArrMult = np.einsum('fij,fjk->fik', ArrMult, Arr[i])
  return ArrMult
def myeinsum(A1, A2):
  return np.einsum('fij,fjk->fik', A1, A2)
A1 = loopMult(Arr)
A2 = reduce(myeinsum, Arr)
A3 = fmuls.multimul(Arr.T).T
print np.allclose(A1,A2)
print np.allclose(A1,A3)
%timeit loopMult(Arr)
%timeit reduce(myeinsum, Arr)
%timeit fmuls.multimul(Arr.T).T

哪些输出

True
True
10 loops, best of 3: 48.4 ms per loop
10 loops, best of 3: 48.8 ms per loop
100 loops, best of 3: 5.82 ms per loop

所以这是一个 8 倍的加速。所有转置的原因是f2py 隐式转置了所有数组，我们需要手动转置它们以告诉它我们的 fortran 代码期望事物被转置。这避免了复制操作。代价是我们的每个 2x2 矩阵都被转置了，所以为了避免执行错误的操作，我们必须反向循环。

超过 8 的加速比应该是可能的 - 我没有花任何时间尝试优化它。

【讨论】：

代码：

返回：