0
我有一些代碼，我正在努力加快與Numba。我已經對這個話題做了一些閱讀，但是我一直無法弄清楚它的100％。Numba不加速功能
下面是代碼：
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功能往往比未即時編譯的人快得多。例如，一旦運行jit，prob_mort的速度會提高大約600％。然而，解算器本身也快不了多少：
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
我知道有些功能不能實時編譯的，所以我換成一個自定義功能JIT的st.binom.pmf功能，實際上減慢時間到每回路大約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
更新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.5倍更快運行。但是，solver_jit（）和solver_norm（）仍然以同樣的速度運行，這意味着在jit函數外部有一些代碼會減慢速度。
您可以發佈您的自定義'binom.pmf'功能？我猜測你使用jit沒有得到任何改進的原因是'intake_dist'在你最內層的循環中，並且這個不能被解決，所以你在求解器中使用了「對象模式」。 – JoshAdel
如果'binom.pmf'是瓶頸，你可以嘗試包裝rmath版本，並通過cffi調用它，正如我在本博文中所描述的：https://www.continuum.io/blog/developer-blog/calling- c-libraries-numba-using-cffi – JoshAdel