代码之家 › 专栏 › 技术社区 › konstant

有没有什么好方法来优化这个python代码的速度?

numba numpy performance arrays python

konstant · 技术社区 · 8 年前

我有下面一段代码,它基本上计算一些数值表达式,并使用它在一定的值范围内进行积分。当前代码在大约 8.6 s ,但我只是使用模拟值,我的实际数组要大得多。尤其是我的实际尺寸 freq_c= (3800, 101) 尺寸为 number_bin = (3800, 100) 这使得下面的代码非常低效,因为实际数组的总执行时间将接近9分钟。代码中相当慢的一部分是对 k_one_third 和 k_two_third ,我也用过 numexpr.evaluate("..") 使代码加速了10-20%。但是,我避免了 numexpr 这样任何人都可以运行它而不必安装包。还有什么方法可以提高这个代码的速度吗?一些因素的改进也就足够了。请注意 for loop 几乎是不可避免的,因为内存问题,因为数组非常大,所以我通过循环一次操纵每个轴。我也想知道 numba jit 优化在这里是可能的。

import numpy as np
import scipy 
from scipy.integrate import simps as simps
import time

def k_one_third(x):
    return (2.*np.exp(-x**2)/x**(1/3) + 4./x**(1/6)*np.exp(-x)/(1+x**(1/3)))**2

def k_two_third(x):
    return (np.exp(-x**2)/x**(2/3) + 2.*x**(5/2)*np.exp(-x)/(6.+x**3))**2

def spectrum(freq_c, number_bin, frequency, gamma, theta):
    theta_gamma_factor = np.einsum('i,j->ij', theta**2, gamma**2)
    theta_gamma_factor += 1.
    t_g_bessel_factor = 1.-1./theta_gamma_factor
    number = np.concatenate((number_bin, np.zeros((number_bin.shape[0], 1), dtype=number_bin.dtype)), axis=1)
    number_theta_gamma = np.einsum('jk, ik->ijk', theta_gamma_factor**2*1./gamma**3, number)
    final = np.zeros((np.size(freq_c[:,0]), np.size(theta), np.size(frequency)))
    for i in xrange(np.size(frequency)):
        b_n_omega_theta_gamma = frequency[i]**2*number_theta_gamma
        eta = theta_gamma_factor**(1.5)*frequency[i]/2.
        eta = np.einsum('jk, ik->ijk', eta, 1./freq_c)
        bessel_eta = np.einsum('jl, ijl->ijl',t_g_bessel_factor, k_one_third(eta))
        bessel_eta += k_two_third(eta)
        eta = None
        integrand = np.multiply(bessel_eta, b_n_omega_theta_gamma, out= bessel_eta)
        final[:,:, i] = simps(integrand, gamma)
        integrand = None
    return final

frequency = np.linspace(1, 100, 100)
theta = np.linspace(1, 3, 100)
gamma = np.linspace(2, 200, 101)
freq_c = np.random.randint(1, 200, size=(50, 101))
number_bin = np.random.randint(1, 100, size=(50, 100))
time1 = time.time()
spectra = spectrum(freq_c, number_bin, frequency, gamma, theta)
print(time.time()-time1)

2 回复 | 直到 8 年前

max9111 8 年前

正如评论中所说,应该重写大部分代码以获得最佳性能。

我只修改了辛普森集成和@hyry答案一点。这加快了计算速度 22.15s 到 1.76s (15倍),根据您提供的测试数据。用简单的循环替换np.einsums,结果应该不到一秒钟。(改进后的集成约0.4s,24秒后 k_one_two_third(x) )

使用numba获得性能 read . 最新的numba版本(0.39)、IntelSVML包以及fastMath=true等对您的示例产生了很大的影响。

代码

#a bit faster than HYRY's version
@nb.njit(parallel=True,fastmath=True,error_model='numpy')
def k_one_two_third(x):
  one=np.empty(x.shape,dtype=x.dtype)
  two=np.empty(x.shape,dtype=x.dtype)
  for i in nb.prange(x.shape[0]):
    for j in range(x.shape[1]):
      for k in range(x.shape[2]):
        x0 = x[i,j,k] ** (1/3)
        x1 = np.exp(-x[i,j,k] ** 2)
        x2 = np.exp(-x[i,j,k])
        one[i,j,k] = (2*x1/x0 + 4*x2/(x[i,j,k]**(1/6)*(x0 + 1)))**2
        two[i,j,k] = (2*x[i,j,k]**(5/2)*x2/(x[i,j,k]**3 + 6) + x1/x[i,j,k]**(2/3))**2
  return one, two

#improved integration
@nb.njit(fastmath=True)
def simpson_nb(y_in,dx):
  s = y[0]+y[-1]

  n=y.shape[0]//2
  for i in range(n-1):
    s += 4.*y[i*2+1]
    s += 2.*y[i*2+2]

  s += 4*y[(n-1)*2+1]
  return(dx/ 3.)*s

@nb.jit(fastmath=True)
def spectrum(freq_c, number_bin, frequency, gamma, theta):
    theta_gamma_factor = np.einsum('i,j->ij', theta**2, gamma**2)
    theta_gamma_factor += 1.
    t_g_bessel_factor = 1.-1./theta_gamma_factor
    number = np.concatenate((number_bin, np.zeros((number_bin.shape[0], 1), dtype=number_bin.dtype)), axis=1)
    number_theta_gamma = np.einsum('jk, ik->ijk', theta_gamma_factor**2*1./gamma**3, number)
    final = np.empty((np.size(frequency),np.size(freq_c[:,0]), np.size(theta)))

    #assume that dx is const. on integration
    #speedimprovement of the scipy.simps is about 4x
    #numba version to scipy.simps(y,x) is about 60x
    dx=gamma[1]-gamma[0]

    for i in range(np.size(frequency)):
        b_n_omega_theta_gamma = frequency[i]**2*number_theta_gamma
        eta = theta_gamma_factor**(1.5)*frequency[i]/2.
        eta = np.einsum('jk, ik->ijk', eta, 1./freq_c)

        one,two=k_one_two_third(eta)

        bessel_eta = np.einsum('jl, ijl->ijl',t_g_bessel_factor, one)
        bessel_eta += two

        integrand = np.multiply(bessel_eta, b_n_omega_theta_gamma, out= bessel_eta)

        #reorder array
        for j in range(integrand.shape[0]):
          for k in range(integrand.shape[1]):
            final[i,j, k] = simpson_nb(integrand[j,k,:],dx)
    return final

HYRY 8 年前

我分析了代码,发现 k_one_third() 和 k_two_third() 很慢。这两个函数中有一些重复的计算。

将两个函数合并为一个函数,并用 @numba.jit(parallel=True) 我加速了4倍。

@jit(parallel=True)
def k_one_two_third(x):
    x0 = x ** (1/3)
    x1 = np.exp(-x ** 2)
    x2 = np.exp(-x)
    one = (2*x1/x0 + 4*x2/(x**(1/6)*(x0 + 1)))**2
    two = (2*x**(5/2)*x2/(x**3 + 6) + x1/x**(2/3))**2
    return one, two