优化一段Julia和Python代码的建议

时间:2019-06-22 21:46:58

标签: python performance optimization julia numba

我有一个用于仿真的程序,并且在程序中有一个功能。我已经意识到该功能会消耗大部分仿真时间。因此,我尝试首先优化该功能。功能如下

Julia 1.1版:

function fun_jul(M,ksi,xi,x)

  F(n,x) = sin(n*pi*(x+1)/2)*cos(n*pi*(x+1)/2);
  K = length(ksi);
  Z = zeros(length(x),K);
  for n in 1:M
    for k in 1:K
        for l in 1:length(x)
          Z[l,k] +=  (1-(n/(M+1))^2)^xi*F(n,ksi[k])*F(n,x[l]);
        end
    end
  end

return Z

end

我还用python + numba重写了上面的函数,以进行比较,如下所示

Python + numba

import numpy as np
from numba import prange, jit    
@jit(nopython=True, parallel=True)
def fun_py(M,ksi,xi,x):

    K = len(ksi);
    F = lambda nn,xx: np.sin(nn*np.pi*(xx+1)/2)*np.cos(nn*np.pi*(xx+1)/2);

    Z = np.zeros((len(x),K));
    for n in range(1,M+1):
        for k in prange(0,K):
            Z[:,k] += (1-(n/(M+1))**2)**xi*F(n,ksi[k])*F(n,x);


    return Z

但是茱莉亚代码很慢,这是我的结果:

朱莉娅成绩:

using BenchmarkTools
N=400; a=-0.5; b=0.5; x=range(a,b,length=N); cc=x; M = 2*N+100; xi = M/40;
@benchmark fun_jul(M,cc,xi,x)

BenchmarkTools.Trial: 
  memory estimate:  1.22 MiB
  allocs estimate:  2
  --------------
  minimum time:     25.039 s (0.00% GC)
  median time:      25.039 s (0.00% GC)
  mean time:        25.039 s (0.00% GC)
  maximum time:     25.039 s (0.00% GC)
  --------------
  samples:          1
  evals/sample:     1

Python结果:

N=400;a = -0.5;b = 0.5;x = np.linspace(a,b,N);cc = x;M = 2*N + 100;xi = M/40;

%timeit fun_py(M,cc,xi,x);
1.2 s ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

对于改善julia和python + numba的代码的任何帮助,将不胜感激。

已更新

基于@Przemyslaw Szufel的回答和其他帖子,我改进了numba和julia代码。现在两者都已并行化。这是时间

Python + Numba时间:

@jit(nopython=True, parallel=True)
def fun_py(M,ksi,xi,x):

    K = len(ksi);
    F = lambda nn,xx: np.sin(nn*np.pi*(xx+1)/2)*np.cos(nn*np.pi*(xx+1)/2);

Z = np.zeros((K,len(x)));
for n in range(1,M+1):
    pw = (1-(n/(M+1))**2)**xi; f=F(n,x)
    for k in prange(0,K):
        Z[k,:] = Z[k,:] + pw*F(n,ksi[k])*f;


    return Z


N=1000; a=-0.5; b=0.5; x=np.linspace(a,b,N); cc=x; M = 2*N+100; xi = M/40;

%timeit fun_py(M,cc,xi,x);
733 ms ± 13.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

朱莉亚时代

N=1000; a=-0.5; b=0.5; x=range(a,b,length=N); cc=x; M = 2*N+100; xi = M/40;
@benchmark fun_jul2(M,cc,xi,x)

BenchmarkTools.Trial: 
  memory estimate:  40.31 MiB
  allocs estimate:  6302
  --------------
  minimum time:     705.470 ms (0.17% GC)
  median time:      726.403 ms (0.17% GC)
  mean time:        729.032 ms (1.68% GC)
  maximum time:     765.426 ms (5.27% GC)
  --------------
  samples:          7
  evals/sample:     1

3 个答案:

答案 0 :(得分:5)

我使用以下代码在单线程上(而不是在我的计算机上为28s)下降到300ms。

您正在为Numba使用多线程。在Julia中,您应该使用并行处理(Julia对多线程支持是实验性的)。看来您的代码正在进行某种参数扫描-这样的代码很容易并行化,但通常需要对您的计算过程进行一些调整。

代码如下:

function fun_jul2(M,ksi,xi,x)
  F(n,x) = sin(n*pi*(x+1))/2;
  K = length(ksi);
  L = length(x);
  Z = zeros(length(x),K);
  for n in 1:M
    F_im1= [F(n,ksi[k]) for k in 1:K]
    F_im2 = [F(n,x[l]) for l in 1:L]
    pow = (1-(n/(M+1))^2)^xi    
    for k in 1:K
       for l in 1:L    
          Z[l,k] += pow*F_im1[k]*F_im2[l];
      end
    end
  end
  Z
end
julia> fun_jul2(M,cc,xi,x) ≈ fun_jul(M,cc,xi,x)
true

julia> @time fun_jul2(M,cc,xi,x);
  0.305269 seconds (1.81 k allocations: 6.934 MiB, 1.60% gc time)

**编辑:具有DNF建议的多线程和入站:

function fun_jul3(M,ksi,xi,x)
  F(n,x) = sin(n*pi*(x+1))/2;
  K = length(ksi);
  L = length(x);
  Z = zeros(length(x),K);
  for n in 1:M
    F_im1= [F(n,ksi[k]) for k in 1:K]
    F_im2 = [F(n,x[l]) for l in 1:L]
    pow = (1-(n/(M+1))^2)^xi    
    Threads.@threads for k in 1:K
       for l in 1:L    
          @inbounds Z[l,k] += pow*F_im1[k]*F_im2[l];
      end
    end
  end
  Z
end

现在是运行时间(在启动Julia之前,请记住运行set JULIA_NUM_THREADS=4或Linux等效版本):

julia>  fun_jul2(M,cc,xi,x) ≈ fun_jul3(M,cc,xi,x)
true

julia> @time fun_jul3(M,cc,xi,x);
  0.051470 seconds (2.71 k allocations: 6.989 MiB)

您还可以尝试进一步并行化F_im1F_im2的计算。

答案 1 :(得分:2)

您可以使用任何具有循环的语言来执行或不执行循环优化。此处的主要区别在于,numba代码是针对内部循环进行矢量化的,而Julia代码不是。要对Julia版本进行矢量化处理,有时需要使用。将运算符更改为其矢量化版本,例如+变为。+。

由于我无法让Numba正确安装在较旧的Windows 10计算机上,因此我在Web上的免费Linux版本上运行了以下代码版本。这意味着我不得不将Python接口用于timeit(),而不是命令行。

在mybinder的Jupyter中运行,可能没有指定线程,因此可能只有1个线程。 :

import timeit
timeit.timeit("""
@jit(nopython=True, parallel=True)
def fun_py(M,ksi,xi,x):

    K = len(ksi);
    F = lambda nn,xx: np.sin(nn*np.pi*(xx+1)/2)*np.cos(nn*np.pi*(xx+1)/2);

    Z = np.zeros((len(x),K));
    for n in range(1,M+1):
        for k in prange(0,K):
            Z[:,k] += (1-(n/(M+1))**2)**xi*F(n,ksi[k])*F(n,x);


    return Z


N=400; a = -0.5; b = 0.5; x = np.linspace(a,b,N); cc = x;M = 2*N + 100; xi = M/40;
fun_py(M,cc,xi,x)
""", setup ="import numpy as np; from numba import prange, jit", number=5)

出[1]:61.07768889795989

您的计算机必须比Jupyter,ForBonder快得多。

下面我在JuliaBox的Jupyter中运行了此优化的julia代码版本,指定了1个线程内核:

using BenchmarkTools

F(n, x) = sinpi.(n * (x .+ 1) / 2) .* cospi.(n * (x .+ 1) / 2)

function fun_jul2(M, ksi, xi, x)
    K = length(ksi)
    Z = zeros(length(x), K)
    for n in 1:M, k in 1:K
        Z[:, k] .+= (1 - (n / (M + 1))^2)^xi * F(n, ksi[k]) * F(n, x)
    end
    return Z
end


const N=400; const a=-0.5; const b=0.5; const x=range(a,b,length=N); 
const cc=x; const M = 2*N+100; const xi = M/40;

@btime fun_jul2(M, cc, xi, x)

8.076 s(1080002分配:3.35 GiB)

答案 2 :(得分:2)

为获得性能,只需预先计算三角部分。 实际上,sin是一项昂贵的操作:

%timeit np.sin(1.)
712 ns ± 2.22 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit 1.2*3.4
5.88 ns ± 0.016 ns per loop (mean ± std. dev. of 7 runs, 100000000 loops each)

在python中:

@jit
def fun_py2(M,ksi,xi,x):

    NN = np.arange(1,M+1)
    Fksi = np.sin(np.pi*np.outer(NN,ksi+1))/2  # sin(a)cos(a) is sin(2a)/2
    Fx = np.sin(np.pi*np.outer(NN,x+1))/2
    U = (1-(NN/(M+1))**2)**xi
    Z = np.zeros((len(x),len(ksi)))

    for n in range(len(NN)):
        for k in range(len(ksi)):
            for l in range(len(x)):
                Z[k,l] += U[n] * Fksi[n,k] * Fx[n,l];
    return Z

要提高30倍:

np.allclose(fun_py(M,cc,xi,x),fun_py2(M,cc,xi,x))
True

%timeit fun_py(M,cc,xi,x)
1.14 s ± 4.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit fun_py2(M,cc,xi,x)
29.5 ms ± 375 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

这不会触发任何并行性。我想朱莉娅也会发生同样的事情。