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