python - tutorial - ¿De dónde vienen las ganancias usando numba para código numpy puro?

matplotlib title position

Question

Me gustaría entender de dónde provienen las ganancias cuando se usa Numba para acelerar el código numpy puro en un bucle for. ¿Hay alguna herramienta de creación de perfiles que le permita ver las funciones de jitted ?

El código de demostración (como se muestra a continuación) solo utiliza una multiplicación de matrices muy básica para proporcionar trabajo a la computadora. Son las ganancias observadas de:

1. un loop más rápido,
2. la refundición de numpy funciones numpy interceptadas por el jit durante el proceso de compilación, o
3. menos sobrecarga con jit, ya que la ejecución de subcontratación numpy se realiza a través de funciones de envoltorio a bibliotecas de bajo nivel como LINPACK

%matplotlib inline
import numpy as np
from numba import jit
import pandas as pd
#Dimensions of Matrices
i = 100
j = 100
def pure_python(N,i,j):
for n in range(N):
a = np.random.rand(i,j)
b = np.random.rand(i,j)
c = np.dot(a,b)
@jit(nopython=True)
def jit_python(N,i,j):
for n in range(N):
a = np.random.rand(i,j)
b = np.random.rand(i,j)
c = np.dot(a,b)
time_python = []
time_jit = []
N = [1,10,100,500,1000,2000]
for n in N:
time = %timeit -oq pure_python(n,i,j)
time_python.append(time.average)
time = %timeit -oq jit_python(n,i,j)
time_jit.append(time.average)
df = pd.DataFrame({''pure_python'' : time_python, ''jit_python'' : time_jit}, index=N)
df.index.name = ''Iterations''
df[["pure_python", "jit_python"]].plot()

Produce la siguiente tabla.

Answer 1

TL: DR El aleatorio y el bucle se aceleran, pero la matriz no lo hace excepto para un tamaño de matriz pequeño. Con un tamaño de matriz / bucle pequeño, parece que hay aceleraciones significativas que probablemente están relacionadas con la sobrecarga de Python. Con N grande, la matriz comienza a dominar y la jit es menos útil.

Definiciones de funciones, utilizando una matriz cuadrada por simplicidad.

from IPython.display import display
import numpy as np
from numba import jit
import pandas as pd
#Dimensions of Matrices
N = 1000
def py_rand(i, j):
a = np.random.rand(i, j)
jit_rand = jit(nopython=True)(py_rand)
def py_matmul(a, b):
c = np.dot(a, b)
jit_matmul = jit(nopython=True)(py_matmul)
def py_loop(N, val):
count = 0
for i in range(N):
count += val
jit_loop = jit(nopython=True)(py_loop)
def pure_python(N,i,j):
for n in range(N):
a = np.random.rand(i,j)
b = np.random.rand(i,j)
c = np.dot(a,a)
jit_func = jit(nopython=True)(pure_python)

Sincronización:

df = pd.DataFrame(columns=[''Func'', ''jit'', ''N'', ''Time''])
def meantime(f, *args, **kwargs):
t = %timeit -oq -n5 f(*args, **kwargs)
return t.average
for N in [10, 100, 1000, 2000]:
a = np.random.randn(N, N)
b = np.random.randn(N, N)
df = df.append({''Func'': ''jit_rand'', ''N'': N, ''Time'': meantime(jit_rand, N, N)}, ignore_index=True)
df = df.append({''Func'': ''py_rand'', ''N'': N, ''Time'': meantime(py_rand, N, N)}, ignore_index=True)
df = df.append({''Func'': ''jit_matmul'', ''N'': N, ''Time'': meantime(jit_matmul, a, b)}, ignore_index=True)
df = df.append({''Func'': ''py_matmul'', ''N'': N, ''Time'': meantime(py_matmul, a, b)}, ignore_index=True)
df = df.append({''Func'': ''jit_loop'', ''N'': N, ''Time'': meantime(jit_loop, N, 2.0)}, ignore_index=True)
df = df.append({''Func'': ''py_loop'', ''N'': N, ''Time'': meantime(py_loop, N, 2.0)}, ignore_index=True)
df = df.append({''Func'': ''jit_func'', ''N'': N, ''Time'': meantime(jit_func, 5, N, N)}, ignore_index=True)
df = df.append({''Func'': ''py_func'', ''N'': N, ''Time'': meantime(pure_python, 5, N, N)}, ignore_index=True)
df[''jit''] = df[''Func''].str.contains(''jit'')
df[''Func''] = df[''Func''].apply(lambda s: s.split(''_'')[1])
df.set_index(''Func'')
display(df)

resultado:

 Func jit N Time
0 rand True 10 1.030686e-06
1 rand False 10 1.115149e-05
2 matmul True 10 2.250371e-06
3 matmul False 10 2.199343e-06
4 loop True 10 2.706000e-07
5 loop False 10 7.274286e-07
6 func True 10 1.217046e-05
7 func False 10 2.495837e-05
8 rand True 100 5.199217e-05
9 rand False 100 8.149794e-05
10 matmul True 100 7.848071e-05
11 matmul False 100 2.130794e-05
12 loop True 100 2.728571e-07
13 loop False 100 3.003743e-06
14 func True 100 6.739634e-04
15 func False 100 1.146594e-03
16 rand True 1000 5.644258e-03
17 rand False 1000 8.012790e-03
18 matmul True 1000 1.476098e-02
19 matmul False 1000 1.613211e-02
20 loop True 1000 2.846572e-07
21 loop False 1000 3.539849e-05
22 func True 1000 1.256926e-01
23 func False 1000 1.581177e-01
24 rand True 2000 2.061612e-02
25 rand False 2000 3.204709e-02
26 matmul True 2000 9.866484e-02
27 matmul False 2000 1.007234e-01
28 loop True 2000 3.011143e-07
29 loop False 2000 7.477454e-05
30 func True 2000 1.033560e+00
31 func False 2000 1.199969e+00

Parece que numba está optimizando el bucle, así que no me molestaré en incluirlo en la comparación.

trama:

def jit_speedup(d):
py_time = d[d[''jit''] == False][''Time''].mean()
jit_time = d[d[''jit''] == True][''Time''].mean()
return py_time / jit_time
import seaborn as sns
result = df.groupby([''Func'', ''N'']).apply(jit_speedup).reset_index().rename(columns={0: ''Jit Speedup''})
result = result[result[''Func''] != ''loop'']
sns.factorplot(data=result, x=''N'', y=''Jit Speedup'', hue=''Func'')

Entonces, dado que el bucle es de 5 repeticiones, el jit acelera las cosas bastante sólidamente hasta que la matriz se vuelve lo suficientemente cara como para hacer que la otra sobrecarga sea insignificante en comparación.