Python, numpy, einsum умножают стек матриц

По соображениям производительности

Мне любопытно, есть ли способ умножить стек стека матриц. У меня есть 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 итераций через цикл.

Ответы

Ответ 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). Затем вы будете использовать его в своем script как

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 неявно переносит все массивы, и нам нужно их вручную транспонировать, чтобы сказать, что наш фортран-код ожидает, что вещи будут транспонированы. Это позволяет избежать операции копирования. Стоимость состоит в том, что каждая из наших матриц 2x2 транспонируется, поэтому, чтобы избежать неправильной операции, мы должны выполнить обратное преобразование.

Большее ускорение, чем 8, должно быть возможным - я не тратил времени, пытаясь оптимизировать это.