Ottimizzare doppio loop in pitone

Question

Sto cercando di ottimizzare il ciclo seguente:

def numpy(nx, nz, c, rho): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3): 
      a[ix, iz] = sum(c*rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum(c*rho[ix-2:ix+2, iz]) 
    return a, b 

ho provato diverse soluzioni e trovato usando numba per calcolare la somma del prodotto porta a prestazioni migliori:

import numpy as np 
import numba as nb 
import time 

@nb.autojit 
def sum_opt(arr1, arr2): 
    s = arr1[0]*arr2[0] 
    for i in range(1, len(arr1)): 
     s+=arr1[i]*arr2[i] 
    return s 

def numba1(nx, nz, c, rho): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return a, b 

@nb.autojit 
def numba2(nx, nz, c, rho): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return a, b 
nx = 1024 
nz = 256  

rho = np.random.rand(nx, nz) 
c = np.random.rand(4) 
a = np.zeros((nx, nz)) 
b = np.zeros((nx, nz)) 

ti = time.clock() 
a, b = numpy(nx, nz, c, rho) 
print 'Time numpy : ' + `round(time.clock() - ti, 4)` 

ti = time.clock() 
a, b = numba1(nx, nz, c, rho) 
print 'Time numba1 : ' + `round(time.clock() - ti, 4)` 

ti = time.clock() 
a, b = numba2(nx, nz, c, rho) 
print 'Time numba2 : ' + `round(time.clock() - ti, 4)` 

Questo ha portato ad

Tempo NumPy: 4,1595

Tempo numba1: 0,6993

Tempo numba2: 1,0135

Utilizzando la versione numba della funzione sum (sum_opt) funziona molto bene. Ma mi chiedo perché la versione numba della funzione double loop (numba2) porti a tempi di esecuzione più lenti. Ho provato a usare jit invece di autojit, specificando i tipi di argomento, ma era peggio.

Ho anche notato che il looping sul loop più piccolo è più lento del looping sul loop più grande. C'è qualche spiegazione?

Sia che sia, sono sicuro che questa funzione doppio ciclo può essere migliorata molto vettorizzazione il problema (come this) o utilizzando un altro metodo (mappa?), Ma io sono un po 'confuso su questi metodi.

Nelle altre parti del mio codice, ho usato metodi di affettamento numba e numpy per sostituire tutti i cicli espliciti, ma in questo caso particolare, non so come configurarlo.

Qualche idea?

EDIT

Grazie per tutti i vostri commenti. Ho lavorato un po 'su questo problema:

import numba as nb 
import numpy as np 
from scipy import signal 
import time 


@nb.jit(['float64(float64[:], float64[:])'], nopython=True) 
def sum_opt(arr1, arr2): 
    s = arr1[0]*arr2[0] 
    for i in xrange(1, len(arr1)): 
     s+=arr1[i]*arr2[i] 
    return s 

@nb.autojit 
def numba1(nx, nz, c, rho, a, b): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return a, b 


@nb.jit(nopython=True) 
def numba2(nx, nz, c, rho, a, b): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return a, b 

@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True) 
def numba3a(nx, nz, c, rho, a): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
    return a 

@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True) 
def numba3b(nx, nz, c, rho, b): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return b 

def convol(nx, nz, c, aa, bb): 
    s1 = rho[1:nx-1,2:nz-3] 
    s2 = rho[0:nx-2,2:nz-3] 
    kernel = c[:,None][::-1] 
    aa[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid') 
    bb[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid') 
    return aa, bb 


nx = 1024 
nz = 256 
rho = np.random.rand(nx, nz) 
c = np.random.rand(4) 
a = np.zeros((nx, nz)) 
b = np.zeros((nx, nz)) 

ti = time.clock() 
for i in range(1000): 
    a, b = numba1(nx, nz, c, rho, a, b) 
print 'Time numba1 : ' + `round(time.clock() - ti, 4)` 

ti = time.clock() 
for i in range(1000): 
    a, b = numba2(nx, nz, c, rho, a, b) 
print 'Time numba2 : ' + `round(time.clock() - ti, 4)` 

ti = time.clock() 
for i in range(1000): 
    a = numba3a(nx, nz, c, rho, a) 
    b = numba3b(nx, nz, c, rho, b) 
print 'Time numba3 : ' + `round(time.clock() - ti, 4)` 

ti = time.clock() 
for i in range(1000): 
    a, b = convol(nx, nz, c, a, b) 
print 'Time convol : ' + `round(time.clock() - ti, 4)` 

La soluzione è molto elegante Divakar, ma devo usare questa funzione di un gran numero di tempo nel mio codice. Così, per 1000 iterazioni, questo cavo al

Tempo numba1: 3,2487

Tempo numba2: 3,7012

Tempo numba3: 3,2088

Tempo convol: 22,7696

autojit e jit sono molto vicini. Tuttavia, quando si utilizza jit, sembra importante specificare tutti i tipi di argomenti.

Non so se esiste un modo per specificare i tipi di argomento nel decoratore jit quando la funzione ha più uscite. Qualcuno ?

Per ora non ho trovato altra soluzione che usare numba. Nuove idee sono benvenute!

Answer 1

Numba è molto veloce in nopython mode ma con il tuo codice deve tornare alla modalità object, che è molto più lenta. Puoi vedere questo accadendo se passi nopython=True al decoratore jit.

lo fa compilare in nopython modalità (almeno nella versione 0.18.2 Numba) se si passa a e b come argomenti:

import numba as nb 

@nb.jit(nopython=True) 
def sum_opt(arr1, arr2): 
    s = arr1[0]*arr2[0] 
    for i in range(1, len(arr1)): 
     s+=arr1[i]*arr2[i] 
    return s 

@nb.jit(nopython=True) 
def numba2(nx, nz, c, rho, a, b): 
    for ix in range(2, nx-3): 
     for iz in range(2, nz-3):   
      a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz]) 
      b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz]) 
    return a, b 

Si noti che nel release notes si parla di autojit essere sconsigliato a favore di jit.

---

A quanto pare non sei ancora soddisfatto. Che ne dici di una soluzione basata su stride_tricks?

from numpy.lib.stride_tricks import as_strided 

def stridetrick_einsum(c, rho, out): 
    ws = len(c) 
    nx, nz = rho.shape 

    shape = (nx-ws+1, ws, nz) 
    strides = (rho.strides[0],) + rho.strides 
    rho_windowed = as_strided(rho, shape, strides) 

    np.einsum('j,ijk->ik', c, rho_windowed, out=out) 

a = np.zeros((nx, nz)) 
stridetrick_einsum(c, rho[1:-1,2:-3], a[2:-3,2:-3]) 
b = np.zeros((nx, nz)) 
stridetrick_einsum(c, rho[0:-2,2:-3], b[2:-3,2:-3]) 

Cosa c'è di più, dal momento che a e b sono ovviamente quasi esattamente la stessa, è possibile calcolare in una volta sola e poi copiare i valori sopra:

a = np.zeros((nx, nz)) 
stridetrick_einsum(c, rho[:-1,2:-3], a[1:-3,2:-3]) 
b = np.zeros((nx, nz)) 
b[2:-3,2:-3] = a[1:-4,2:-3] 
a[1,2:-3] = 0.0 

Answer 2

Come indicato nella performance section sul blog di continuità, autojit compila just-in-time, mentre jit compila prima del tempo:

Numba può compilare just-in-time con il decoratore autojit, o in anticipo tempo con il decoratore JIT

Ciò significa che, in molti casi, autojit significa che il compilatore può fare una congettura più istruiti per il codice si compila, e ottimizzare in seguito. So che la compilazione just-in-time prima del tempo sembra contraddittoria, ma hey.

ma mi chiedo il motivo per cui la versione numba della funzione doppio anello (numba2) porta a tempi di esecuzione più lenti

Numba non aumenta le prestazioni delle chiamate di funzione arbitrarie. Anche se non posso dirlo con certezza, la mia ipotesi è che il sovraccarico della raccolta JIT superi il beneficio derivante dal farlo (se c'è qualche vantaggio).

Ho anche notato che il ciclo prima sul loop più piccolo è più lento del ciclo prima sul loop più grande. C'è qualche spiegazione?

Ciò è probabilmente dovuto a un cache miss. Un array bidimensionale viene allocato come un blocco contiguo di memoria della dimensione rows * columns. Ciò che viene recuperato nella cache si basa su una combinazione di località temporale (ciò che è stato recentemente usato) e spaziale (ciò che è vicino alla memoria a ciò che è stato usato), cioè ciò che si crede di essere usato in seguito.

In caso di iterazione su più righe, si scorre nell'ordine in cui i dati vengono visualizzati in memoria. Quando si esegue il iter sulle colonne per primo, si "salta" la larghezza di una riga in memoria ogni volta, rendendo meno probabile che la posizione di memoria a cui si accede sia stata scaricata nella cache.

2D array: [[1,2,3],[4,5,6],[7,8,9]] 
In memory: 1 2 3 4 5 6 7 8 9 

Supponiamo un algoritmo di raccolta cache stupido, troppo semplificato, che recupera 3 posizioni di memoria successive.

Iterazione fila-prima:

In memory: 1 2 3 | 4 5 6 | 7 8 9 
Accessed:  1 2 3 | 4 5 6 | 7 8 9 
Cache miss: - - - | - - - | - - - 

Iterazione colonna-prima:

In memory: 1 2 3 | 4 5 6 | 7 8 9 
Accessed:  1 4 7 | 2 5 8 | 3 6 9 
Cache miss: - - - | x x x | x x x 

Answer 3

Lei non è di sfruttare appieno le capacità di numpy. Il numpythonic modo di gestire il problema sarebbe qualcosa di simile:

cs = np.zeros((nx+1, nz)) 
np.cumsum(c*rho, axis=0, out=cs[1:]) 
aa = cs[5:, 2:-3] - cs[1:-4, 2:-3] 
bb = cs[4:-1, 2:-3] - cs[:-5, 2:-3] 

aa sarà ora tenere la, non-zero parte centrale della vostra a array:

>>> a[:5, :5] 
array([[ 0.  , 0.  , 0.  , 0.  , 0.  ], 
     [ 0.  , 0.  , 0.  , 0.  , 0.  ], 
     [ 0.  , 0.  , 2.31296595, 2.15743042, 2.5853117 ], 
     [ 0.  , 0.  , 2.02697233, 2.83191016, 2.58819583], 
     [ 0.  , 0.  , 2.4086584 , 2.45175615, 2.19628507]]) 
>>>aa[:3, :3] 
array([[ 2.31296595, 2.15743042, 2.5853117 ], 
     [ 2.02697233, 2.83191016, 2.58819583], 
     [ 2.4086584 , 2.45175615, 2.19628507]]) 

e allo stesso modo per bb e b .

Sul mio sistema, con l'input di esempio, questo codice viene eseguito oltre 300 volte più veloce della funzione numpy. Secondo i tuoi tempi, questo sarà uno o due ordini di grandezza più veloce di numba.

Answer 4

In questo caso si esegue essenzialmente la convoluzione 2D, con una piccola modifica che il kernel non sta eseguendo l'inversione come fa la normale operazione convolution. Quindi, in sostanza, ci sono due cose che dobbiamo fare qui per utilizzare signal.convolve2d per risolvere il nostro caso -

• Affettare la matrice di ingresso rho per selezionare una parte di esso, che viene utilizzato nella versione originale loopy del codice . Questo sarebbe il dato di input alla convoluzione.
• Invertire il kernel, c e alimentarlo con i dati a fette a signal.convolve2d.

Si prega di notare che questi devono essere fatti per il calcolo di entrambi a e b separatamente.

Ecco l'implementazione -

import numpy as np 
from scipy import signal 

# Slices for convolutions to get a and b respectively   
s1 = rho[1:nx-1,2:nz-3] 
s2 = rho[0:nx-2,2:nz-3] 
kernel = c[:,None][::-1] # convolution kernel 

# Setup output arrays and fill them with convolution results 
a = np.zeros((nx, nz)) 
b = np.zeros((nx, nz)) 

a[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid') 
b[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid') 

Se non avete bisogno gli zeri in più intorno ai confini delle matrici di uscita, si può semplicemente utilizzare le uscite da signal.convolve2d come sono, che deve ulteriormente aumentare il prestazione.

Runtime mette alla prova

In [532]: %timeit loop_based(nx, nz, c, rho) 
1 loops, best of 3: 1.52 s per loop 

In [533]: %timeit numba1(nx, nz, c, rho) 
1 loops, best of 3: 282 ms per loop 

In [534]: %timeit numba2(nx, nz, c, rho) 
1 loops, best of 3: 509 ms per loop 

In [535]: %timeit conv_based(nx, nz, c, rho) 
10 loops, best of 3: 15.5 ms per loop 

Quindi, per il datasize ingresso effettivo, l'approccio basato la circonvoluzione proposto è di circa 100x più velocemente di quanto il codice loopy e circa 20x migliore rispetto al più veloce approccio basato numbanumba1.