Numba не ускоряется функция

Question

У меня есть код, который я пытаюсь ускорить с numba. Я проделал некоторое чтение по этой теме, но я не смог понять это на 100%.

Вот код:

import pandas as pd 
import matplotlib.pyplot as plt 
import numpy as np 
import scipy.stats as st 
import seaborn as sns 
from numba import jit, vectorize, float64, autojit 
sns.set(context='talk', style='ticks', font_scale=1.2, rc={'figure.figsize': (6.5, 5.5), 'xtick.direction': 'in', 'ytick.direction': 'in'}) 

#%% constraints 
x_min = 0        # death below this 
x_max = 20        # maximum weight 
t_max = 100        # maximum time 
foraging_efficiencies = np.linspace(0, 1, 10)    # potential foraging efficiencies 
R = 10.0         # Resource level 

#%% make the body size and time categories 
body_sizes = np.arange(x_min, x_max+1) 
time_steps = np.arange(t_max) 

#%% parameter functions 
@jit 
def metabolic_fmr(x, u,temp):       # metabolic cost function 
    fmr = 0.125*(2**(0.2*temp))*(1 + 0.5*u) + x*0.1 
    return fmr 

def intake_dist(u):       # intake stochastic function (returns a vector) 
    g = st.binom.pmf(np.arange(R+1), R, u) 
    return g 

@jit 
def mass_gain(x, u, temp):      # mass gain function (returns a vector) 
    x_prime = x - metabolic_fmr(x, u,temp) + np.arange(R+1) 
    x_prime = np.minimum(x_prime, x_max) 
    x_prime = np.maximum(x_prime, 0) 
    return x_prime 

@jit 
def prob_attack(P):       # probability of an attack 
    p_a = 0.02*P 
    return p_a 

@jit 
def prob_see(u):       # probability of not seeing an attack 
    p_s = 1-(1-u)**0.3 
    return p_s 

@jit 
def prob_lethal(x):       # probability of lethality given a successful attack 
    p_l = 0.5*np.exp(-0.05*x) 
    return p_l 

@jit 
def prob_mort(P, u, x): 
    p_m = prob_attack(P)*prob_see(u)*prob_lethal(x) 
    return np.minimum(p_m, 1) 

#%% terminal fitness function 
@jit 
def terminal_fitness(x): 
    t_f = 15.0*x/(x+5.0) 
    return t_f 

#%% linear interpolation function 
@jit 
def linear_interpolation(x, F, t): 
    floor = x.astype(int) 
    delta_c = x-floor 
    ceiling = floor + 1 
    ceiling[ceiling>x_max] = x_max 
    floor[floor<x_min] = x_min 
    interpolated_F = (1-delta_c)*F[floor,t] + (delta_c)*F[ceiling,t] 
    return interpolated_F 

#%% solver 
@jit 
def solver_jit(P, temp): 
    F = np.zeros((len(body_sizes), len(time_steps)))   # Expected fitness 
    F[:,-1] = terminal_fitness(body_sizes)    # expected terminal fitness for every body size 
    V = np.zeros((len(foraging_efficiencies), len(body_sizes), len(time_steps)))  # Fitness for each foraging effort 
    D = np.zeros((len(body_sizes), len(time_steps)))   # Decision 
    for t in range(t_max-1)[::-1]: 
     for x in range(x_min+1, x_max+1):    # iterate over every body size except dead 
      for i in range(len(foraging_efficiencies)):  # iterate over every possible foraging efficiency 
       u = foraging_efficiencies[i] 
       g_u = intake_dist(u)    # calculate the distribution of intakes 
       xp = mass_gain(x, u, temp)   # calculate the mass gain 
       p_m = prob_mort(P, u, x)   # probability of mortality 
       V[i,x,t] = (1 - p_m)*(linear_interpolation(xp, F, t+1)*g_u).sum()  # Fitness calculation 
      vmax = V[:,x,t].max() 
      idx = np.argwhere(V[:,x,t]==vmax).min() 
      D[x,t] = foraging_efficiencies[idx] 
      F[x,t] = vmax 
    return D, F 

def solver_norm(P, temp): 
    F = np.zeros((len(body_sizes), len(time_steps)))   # Expected fitness 
    F[:,-1] = terminal_fitness(body_sizes)    # expected terminal fitness for every body size 
    V = np.zeros((len(foraging_efficiencies), len(body_sizes), len(time_steps)))  # Fitness for each foraging effort 
    D = np.zeros((len(body_sizes), len(time_steps)))   # Decision 
    for t in range(t_max-1)[::-1]: 
     for x in range(x_min+1, x_max+1):    # iterate over every body size except dead 
      for i in range(len(foraging_efficiencies)):  # iterate over every possible foraging efficiency 
       u = foraging_efficiencies[i] 
       g_u = intake_dist(u)    # calculate the distribution of intakes 
       xp = mass_gain(x, u, temp)   # calculate the mass gain 
       p_m = prob_mort(P, u, x)   # probability of mortality 
       V[i,x,t] = (1 - p_m)*(linear_interpolation(xp, F, t+1)*g_u).sum()  # Fitness calculation 
      vmax = V[:,x,t].max() 
      idx = np.argwhere(V[:,x,t]==vmax).min() 
      D[x,t] = foraging_efficiencies[idx] 
      F[x,t] = vmax 
    return D, F 

Отдельные JIT функции, как правило, гораздо быстрее, чем не-JITted них. Например, prob_mort примерно на 600% быстрее, когда он запускается через jit. Однако сам решатель не намного быстрее:

In [3]: %timeit -n 10 solver_jit(200, 25) 
10 loops, best of 3: 3.94 s per loop 

In [4]: %timeit -n 10 solver_norm(200, 25) 
10 loops, best of 3: 4.09 s per loop 

Я знаю, что некоторые функции не могут быть JITted, поэтому я заменил функцию st.binom.pmf с пользовательской функцией JIT и что на самом деле замедлилось время примерно до 17 с за цикл, более чем в 5 раз медленнее. Предположительно, потому что scipy функции, на данный момент, сильно оптимизированы.

Поэтому я подозреваю, что медлительность находится либо в функции linear_interpolate, либо где-то в коде решателя за пределами jitted-функций (потому что в какой-то момент я отключил все функции и запустил solver_norm и получил то же время). Любые мысли о том, где будет медленная часть и как ее ускорить?

UPDATE

Вот бином код, который я использовал в попытке ускорить JIT

@jit 
def factorial(n): 
    if n==0: 
     return 1 
    else: 
     return n*factorial(n-1) 

@vectorize([float64(float64,float64,float64)]) 
def binom(k, n, p): 
    binom_coef = factorial(n)/(factorial(k)*factorial(n-k)) 
    pmf = binom_coef*p**k*(1-p)**(n-k) 
    return pmf 

@jit 
def intake_dist(u):       # intake stochastic function (returns a vector) 
    g = binom(np.arange(R+1), R, u) 
    return g 

UPDATE 2 Я попытался запустить мой биномиальный код в nopython режима и узнал, что я делая это неправильно, потому что это было рекурсивным. После фиксации, что путем изменения кода:

@jit(int64(int64), nopython=True) 
def factorial(nn): 
    res = 1 
    for ii in range(2, nn + 1): 
     res *= ii 
    return res 

@vectorize([float64(float64,float64,float64)], nopython=True) 
def binom(k, n, p): 
    binom_coef = factorial(n)/(factorial(k)*factorial(n-k)) 
    pmf = binom_coef*p**k*(1-p)**(n-k) 
    return pmf 

решатель теперь работает на

In [34]: %timeit solver_jit(200, 25) 
1 loop, best of 3: 921 ms per loop 

который составляет около 3.5x быстрее. Однако solver_jit() и solver_norm() все еще работают с одинаковым темпом, а это значит, что за пределами jit-функций существует некоторый код, замедляющий его.

Answer 1

Я смог внести несколько изменений в ваш код, чтобы сделать его таким, чтобы jit-версия могла полностью скомпилироваться в режиме nopython. На моем ноутбуке это результат:

%timeit solver_jit(200, 25) 
1 loop, best of 3: 50.9 ms per loop 

%timeit solver_norm(200, 25) 
1 loop, best of 3: 192 ms per loop 

Для справки я использую Numba 0.27.0. Я признаю, что ошибки компиляции Numba по-прежнему затрудняют определение того, что происходит, но поскольку я играл с ним какое-то время, я создал интуицию для того, что нужно исправлять. Полный код ниже, но вот список изменений, которые я сделал:

• В linear_interpolation изменения x.astype(int) в x.astype(np.int64) поэтому он может составить в режиме nopython.
• В решателе используйте np.sum как функцию, а не метод массива.
• np.argwhere не поддерживается. Напишите настраиваемый цикл.

Возможно, некоторые дополнительные оптимизации могут быть сделаны, но это дает первоначальное ускорение.

Полный код:

import pandas as pd 
import matplotlib.pyplot as plt 
import numpy as np 
import scipy.stats as st 
import seaborn as sns 
from numba import jit, vectorize, float64, autojit, njit 
sns.set(context='talk', style='ticks', font_scale=1.2, rc={'figure.figsize': (6.5, 5.5), 'xtick.direction': 'in', 'ytick.direction': 'in'}) 

#%% constraints 
x_min = 0        # death below this 
x_max = 20        # maximum weight 
t_max = 100        # maximum time 
foraging_efficiencies = np.linspace(0, 1, 10)    # potential foraging efficiencies 
R = 10.0         # Resource level 

#%% make the body size and time categories 
body_sizes = np.arange(x_min, x_max+1) 
time_steps = np.arange(t_max) 

#%% parameter functions 
@njit 
def metabolic_fmr(x, u,temp):       # metabolic cost function 
    fmr = 0.125*(2**(0.2*temp))*(1 + 0.5*u) + x*0.1 
    return fmr 

@njit() 
def factorial(nn): 
    res = 1 
    for ii in range(2, nn + 1): 
     res *= ii 
    return res 

@vectorize([float64(float64,float64,float64)], nopython=True) 
def binom(k, n, p): 
    binom_coef = factorial(n)/(factorial(k)*factorial(n-k)) 
    pmf = binom_coef*p**k*(1-p)**(n-k) 
    return pmf 

@njit 
def intake_dist(u):       # intake stochastic function (returns a vector) 
    g = binom(np.arange(R+1), R, u) 
    return g 

@njit 
def mass_gain(x, u, temp):      # mass gain function (returns a vector) 
    x_prime = x - metabolic_fmr(x, u,temp) + np.arange(R+1) 
    x_prime = np.minimum(x_prime, x_max) 
    x_prime = np.maximum(x_prime, 0) 
    return x_prime 

@njit 
def prob_attack(P):       # probability of an attack 
    p_a = 0.02*P 
    return p_a 

@njit 
def prob_see(u):       # probability of not seeing an attack 
    p_s = 1-(1-u)**0.3 
    return p_s 

@njit 
def prob_lethal(x):       # probability of lethality given a successful attack 
    p_l = 0.5*np.exp(-0.05*x) 
    return p_l 

@njit 
def prob_mort(P, u, x): 
    p_m = prob_attack(P)*prob_see(u)*prob_lethal(x) 
    return np.minimum(p_m, 1) 

#%% terminal fitness function 
@njit 
def terminal_fitness(x): 
    t_f = 15.0*x/(x+5.0) 
    return t_f 

#%% linear interpolation function 
@njit 
def linear_interpolation(x, F, t): 
    floor = x.astype(np.int64) 
    delta_c = x-floor 
    ceiling = floor + 1 
    ceiling[ceiling>x_max] = x_max 
    floor[floor<x_min] = x_min 
    interpolated_F = (1-delta_c)*F[floor,t] + (delta_c)*F[ceiling,t] 
    return interpolated_F 

#%% solver 
@njit 
def solver_jit(P, temp): 
    F = np.zeros((len(body_sizes), len(time_steps)))   # Expected fitness 
    F[:,-1] = terminal_fitness(body_sizes)    # expected terminal fitness for every body size 
    V = np.zeros((len(foraging_efficiencies), len(body_sizes), len(time_steps)))  # Fitness for each foraging effort 
    D = np.zeros((len(body_sizes), len(time_steps)))   # Decision 
    for t in range(t_max-2,-1,-1): 
     for x in range(x_min+1, x_max+1):    # iterate over every body size except dead 
      for i in range(len(foraging_efficiencies)):  # iterate over every possible foraging efficiency 
       u = foraging_efficiencies[i] 
       g_u = intake_dist(u)    # calculate the distribution of intakes 
       xp = mass_gain(x, u, temp)   # calculate the mass gain 
       p_m = prob_mort(P, u, x)   # probability of mortality 
       V[i,x,t] = (1 - p_m)*np.sum((linear_interpolation(xp, F, t+1)*g_u))  # Fitness calculation 
      vmax = V[:,x,t].max() 

      for k in xrange(V.shape[0]): 
       if V[k,x,t] == vmax: 
        idx = k 
        break 
      #idx = np.argwhere(V[:,x,t]==vmax).min() 
      D[x,t] = foraging_efficiencies[idx] 
      F[x,t] = vmax 
    return D, F 

def solver_norm(P, temp): 
    F = np.zeros((len(body_sizes), len(time_steps)))   # Expected fitness 
    F[:,-1] = terminal_fitness(body_sizes)    # expected terminal fitness for every body size 
    V = np.zeros((len(foraging_efficiencies), len(body_sizes), len(time_steps)))  # Fitness for each foraging effort 
    D = np.zeros((len(body_sizes), len(time_steps)))   # Decision 
    for t in range(t_max-1)[::-1]: 
     for x in range(x_min+1, x_max+1):    # iterate over every body size except dead 
      for i in range(len(foraging_efficiencies)):  # iterate over every possible foraging efficiency 
       u = foraging_efficiencies[i] 
       g_u = intake_dist(u)    # calculate the distribution of intakes 
       xp = mass_gain(x, u, temp)   # calculate the mass gain 
       p_m = prob_mort(P, u, x)   # probability of mortality 
       V[i,x,t] = (1 - p_m)*(linear_interpolation(xp, F, t+1)*g_u).sum()  # Fitness calculation 
      vmax = V[:,x,t].max() 
      idx = np.argwhere(V[:,x,t]==vmax).min() 
      D[x,t] = foraging_efficiencies[idx] 
      F[x,t] = vmax 
    return D, F 

Answer 2

Как уже говорилось, существует вероятный код, возвращающийся в режим объекта. Я просто хотел добавить, что вы можете использовать njit вместо jit для отключения режима объекта. Это поможет диагностировать, какой код является виновником.