Ускорение суммирования numpy между большими массивами

Question

У меня есть код работает на numpy массивов. Хотя операции с линейной алгеброй кажутся быстрыми, я теперь обнаруживаю узкое место в другой проблеме: суммирование двух разных массивов. В приведенном ниже примере WE3 и T1 представлены два массива 1000X1000X1000. Сначала я вычисляю WE3, используя операцию numpy, затем суммирую эти массивы.

import numpy as np 
import scipy as sp 
import time 
N = 100 
n = 1000 
X = np.random.uniform(size = (N,n)) 
wE = np.mean(X,0) 
wE3 = np.einsum('i,j,k->ijk', wE, wE, wE) #22 secs 
T1 = np.random.uniform(size = (n,n,n)) 
a = wE3 + T1 #115 secs 

Расчет wE3 занимает, как 22 секунд, в то время как добавление между WE3 и T1 занимает 115 секунд.

Есть ли какая-то известная причина, почему суммирование этих массивов происходит настолько медленнее, чем вычисление WE3? Они должны иметь более или менее такую ​​же сложность.

Есть ли способ ускорить этот код?

Answer 1

Эта часть np.einsum('i,j,k->ijk', wE, wE, wE) не выполняет никаких суммирующих преобразований и, по существу, просто передается по элементарному умножению. Таким образом, мы можем заменить, что-то вроде этого -

wE[:,None,None] * wE[:,None] * wE 

Продолжительность испытания -

In [9]: # Setup inputs at 1/5th of original dataset sizes 
    ...: N = 20 
    ...: n = 200 
    ...: X = np.random.uniform(size = (N,n)) 
    ...: wE = np.mean(X,0) 
    ...: 

In [10]: %timeit np.einsum('i,j,k->ijk', wE, wE, wE) 
10 loops, best of 3: 45.7 ms per loop 

In [11]: %timeit wE[:,None,None] * wE[:,None] * wE 
10 loops, best of 3: 26.1 ms per loop 

---

Далее, мы имеем wE3 + T1, где T1 = np.random.uniform(size = (n,n,n)) не выглядит, как можно было бы помочь в большой путь, поскольку мы должны создать T1 в любом случае, а затем это просто элементарное дополнение. Кажется, мы можем использовать np.add, что позволяет нам записывать результаты в один из массивов: wE3 или T1. Предположим, мы выберем T1, если это нормально для изменения. Я предполагаю, что это принесет небольшую эффективность памяти, поскольку мы не будем добавлять другую переменную в рабочее пространство.

Таким образом, мы могли бы сделать -

np.add(wE3,T1,out=T1) 

время выполнения теста -

In [58]: def func1(wE3): 
    ...:  T1 = np.random.uniform(size = (n,n,n)) 
    ...:  return wE3 + T1 
    ...: 
    ...: def func2(wE3): 
    ...:  T1 = np.random.uniform(size = (n,n,n)) 
    ...:  np.add(wE3,T1,out=T1) 
    ...:  return T1 
    ...: 

In [59]: # Setup inputs at 1/4th of original dataset sizes 
    ...: N = 25 
    ...: n = 250 
    ...: X = np.random.uniform(size = (N,n)) 
    ...: wE = np.mean(X,0) 
    ...: wE3 = np.einsum('i,j,k->ijk', wE, wE, wE) 
    ...: 

In [60]: %timeit func1(wE3) 
1 loops, best of 3: 390 ms per loop 

In [61]: %timeit func2(wE3) 
1 loops, best of 3: 363 ms per loop 

Использование @Aaron's suggestion, мы можем использовать цикл и предполагая, что запись обратно результаты в wE3 в порядке, мы могли бы сделать -

wE3 = wE[:,None,None] * wE[:,None] * wE 
for x in wE3: 
    np.add(x, np.random.uniform(size = (n,n)), out=x) 

Конечные результаты

Таким образом, соединив назад все предложенные усовершенствования, наконец, результаты тестирования выполнения были -

In [97]: def func1(wE): 
    ...:  wE3 = np.einsum('i,j,k->ijk', wE, wE, wE) 
    ...:  T1 = np.random.uniform(size = (n,n,n)) 
    ...:  return wE3 + T1 
    ...: 
    ...: def func2(wE): 
    ...:  wE3 = wE[:,None,None] * wE[:,None] * wE 
    ...:  for x in wE3: 
    ...:   np.add(x, np.random.uniform(size = (n,n)), out=x) 
    ...:  return wE3 
    ...: 

In [98]: # Setup inputs at 1/3rd of original dataset sizes 
    ...: N = 33 
    ...: n = 330 
    ...: X = np.random.uniform(size = (N,n)) 
    ...: wE = np.mean(X,0) 
    ...: 

In [99]: %timeit func1(wE) 
1 loops, best of 3: 1.09 s per loop 

In [100]: %timeit func2(wE) 
1 loops, best of 3: 879 ms per loop 

Answer 2

Есть ли известная причина, почему суммирование этих массивов такое медленнее, чем вычисление WE3?

Массивы wE3, T1 и a каждый требует 8 гигабайт памяти. Вероятно, у вас закончилась физическая память, а доступ к swap memory убивает вашу производительность.

Есть ли способ ускорить этот код?

Получить дополнительную физическую память (то есть ОЗУ).

Если это невозможно, посмотрите, что вы собираетесь делать с этими массивами, и посмотрите, можете ли вы работать партиями, чтобы общая память, необходимая при обработке пакета, оставалась в пределах вашей физической памяти ,

Answer 3

Для этого вам действительно нужно использовать Jem Numba (только во время компилятора). Это чисто бесчисленный конвейер, который идеально подходит для Numba.

Все, что вам нужно сделать, это бросить этот код в функцию и надеть декоратор @jit сверху. Он получает ускорения рядом с Китоном.

Однако, как уже отмечалось, по-видимому, вы пытаетесь работать с данными, слишком большой для вашей локальной машине, и Numba не решит ваши проблемы