朱莉娅 DifferentialEquations.jl 速度

Julia DifferentialEquations.jl speed

我正在尝试测试 Julia ODE 求解器的速度。我在教程中使用了洛伦兹方程:

using DifferentialEquations
using Plots
function lorenz(t,u,du)
du[1] = 10.0*(u[2]-u[1])
du[2] = u[1]*(28.0-u[3]) - u[2]
du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;1.0;1.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz,u0,tspan)
sol = solve(prob,reltol=1e-8,abstol=1e-8,saveat=collect(0:0.01:100))

一开始加载包大约需要 25 秒,代码 运行 在 Jupyter notebook 中的 windows 10 四核笔记本电脑上需要 7 秒。我知道 Julia 需要先预编译包,这就是加载时间这么长的原因吗?我发现 25 秒难以忍受。此外,当我 运行 求解器再次使用不同的初始值时,它花费的时间更少 (~1s) 到 运行,这是为什么?这是典型的速度吗?

Tl;博士:

  1. Julia 包有一个预编译阶段。这有助于使所有进一步的 using 调用更快,但代价是第一个存储一些编译数据。这只会触发每个包更新。
  2. using 必须拉入包,这需要一点时间(取决于可以预编译多少)。
  3. 预编译不是“完整的”,所以你第一次 运行 一个函数,即使是来自一个包,它也必须编译。
  4. Julia 开发人员知道这一点,并且已经计划通过使预编译更完整来摆脱 (2) 和 (3)。还有计划减少编译时间,我不知道具体细节。
  5. 所有 Julia 函数都专注于给定的类型,并且每个函数都是一个单独的类型,因此 DiffEq 的内部函数专注于您提供的每个 ODE 函数。
  6. 在大多数需要长时间计算的情况下,(5) 实际上并不重要,因为您不会经常更改函数(如果经常更改,请考虑更改参数)。
  7. 但是 (6) 在交互式使用时确实很重要。它使它感觉不那么“光滑”。
  8. 我们可以取消对 ODE 函数的这种专门化,但它不是默认设置,因为它会导致 2 到 4 倍的性能损失。说不定以后就是默认了。
  9. 我们在这个问题上的计时 post 预编译仍然比 SciPy 的包装 Fortran 求解器这样的问题好 20 倍。所以这都是编译时间问题,而不是运行时间问题。编译时间基本上是恒定的(调用相同函数的较大问题具有大致相同的编译),所以这实际上只是一个交互性问题。
  10. 我们(和一般的 Julia)将来可以并且会在交互性方面做得更好。

完整说明

这真的不是 DifferentialEquations.jl 的东西,这只是一个 Julia 包的东西。 25s 必须包括预编译时间。第一次加载 Julia 包时,它会进行预编译。然后在下一次更新之前不需要再次发生。这可能是最长的初始化,对于 DifferentialEquations.jl 来说相当长,但同样只会在每次更新包代码时发生。然后,每次 using 都有一个小的初始化成本。 DiffEq 相当大,所以初始化需要一点时间:

@time using DifferentialEquations
5.201393 seconds (4.16 M allocations: 235.883 MiB, 4.09% gc time)

然后如评论中所述,您还有:

@time using Plots
6.499214 seconds (2.48 M allocations: 140.948 MiB, 0.74% gc time)

那么,你第一次运行

function lorenz(t,u,du)
  du[1] = 10.0*(u[2]-u[1])
  du[2] = u[1]*(28.0-u[3]) - u[2]
  du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;1.0;1.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz,u0,tspan)
@time sol = solve(prob,reltol=1e-8,abstol=1e-8,saveat=collect(0:0.01:100))

6.993946 seconds (7.93 M allocations: 436.847 MiB, 1.47% gc time)

然后第二次第三次:

0.010717 seconds (72.21 k allocations: 6.904 MiB)
0.011703 seconds (72.21 k allocations: 6.904 MiB)

这是怎么回事?第一次 Julia 运行 是一个函数时,它会编译它。所以你第一次 运行 solve 时,它会像 运行 那样编译它的所有内部函数。所有进行的时间都将没有编译。 DifferentialEquations.jl也专门针对函数本身,所以如果我们改变函数:

function lorenz2(t,u,du)
  du[1] = 10.0*(u[2]-u[1])
  du[2] = u[1]*(28.0-u[3]) - u[2]
  du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;1.0;1.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz2,u0,tspan)

我们将再次承担一些编译时间:

@time sol = 
solve(prob,reltol=1e-8,abstol=1e-8,saveat=collect(0:0.01:100))
3.690755 seconds (4.36 M allocations: 239.806 MiB, 1.47% gc time)

这是什么,现在是为什么。这里有一些东西。首先,Julia 包不会完全预编译。它们不会在会话之间保留实际方法的缓存编译版本。这是 1.x 发布列表中要做的事情,这将摆脱第一次命中,类似于调用 C/Fortran 包,因为它会提前很多命中(AOT) 编译函数。这样很好,但现在请注意有启动时间。

现在让我们谈谈更改功能。 Julia 中的每个函数都会自动专注于其参数 (see this blog post for details)。这里的关键思想是 Julia 中的每个函数都是一个单独的具体类型。所以,既然这里的问题类型是参数化的,改变函数就会触发编译。请注意它的关系:您可以更改函数的参数(如果您有参数),您可以更改初始条件等,但它只会更改触发重新编译的类型。

值得吗?也许会。我们想专注于快速处理困难的计算。编译时间是恒定的(即您可以解决 6 小时的 ODE,它仍然需要几秒钟),因此计算成本高的计算不会在这里受到影响。 Monte Carlo 您 运行 设置了数千个参数和初始条件的模拟在这里不会受到影响,因为如果您只是更改初始条件和参数的值,那么它不会重新编译。但是在您更改功能的地方进行交互式使用确实会在那里受到第二次左右的打击,这并不好。 Julia 开发人员对此的一个回答是花费 post Julia 1.0 时间来加快编译时间,这是我不知道细节的事情,但我确信这里有一些容易实现的成果。

我们可以摆脱它吗?是的。 DiffEq Online 不会为每个函数重新编译,因为它适合在线使用。

function lorenz3(t,u,du)
  du[1] = 10.0*(u[2]-u[1])
  du[2] = u[1]*(28.0-u[3]) - u[2]
  du[3] = u[1]*u[2] - (8/3)*u[3]
  nothing
end

u0 = [1.0;1.0;1.0]
tspan = (0.0,100.0)
f = NSODEFunction{true}(lorenz3,tspan[1],u0)
prob = ODEProblem{true}(f,u0,tspan)

@time sol = solve(prob,reltol=1e-8,abstol=1e-8,saveat=collect(0:0.01:100))

1.505591 seconds (860.21 k allocations: 38.605 MiB, 0.95% gc time)

现在我们可以更改功能而不产生编译成本:

function lorenz4(t,u,du)
  du[1] = 10.0*(u[2]-u[1])
  du[2] = u[1]*(28.0-u[3]) - u[2]
  du[3] = u[1]*u[2] - (8/3)*u[3]
  nothing
end

u0 = [1.0;1.0;1.0]
tspan = (0.0,100.0)
f = NSODEFunction{true}(lorenz4,tspan[1],u0)
prob = ODEProblem{true}(f,u0,tspan)

@time sol = 
solve(prob,reltol=1e-8,abstol=1e-8,saveat=collect(0:0.01
:100))

0.038276 seconds (242.31 k allocations: 10.797 MiB, 22.50% gc time)

和 tada,通过将函数包装在 NSODEFunction 中(在内部使用 FunctionWrappers.jl) it no longer specializes per-function and you hit the compilation time once per Julia session (and then once that's cached, once per package update). But notice that this has about a 2x-4x cost so I am not sure if it will be enabled by default。我们可以在问题类型构造函数中默认实现这一点(即默认情况下没有额外的特化,但用户可以以交互性为代价选择更快的速度)但我不确定这里更好的默认值是什么(请随意评论你的想法)。但在 Julia 做了它的关键字后,它肯定会很快被记录下来参数更改,因此“免编译”模式将成为使用它的标准方式,即使不是默认方式。

但从长远来看,

import numpy as np
from scipy.integrate import odeint
y0 = [1.0,1.0,1.0]
t = np.linspace(0, 100, 10001)
def f(u,t):
    return [10.0*(u[1]-u[0]),u[0]*(28.0-u[2])-u[1],u[0]*u[1]-(8/3)*u[2]]
%timeit odeint(f,y0,t,atol=1e-8,rtol=1e-8)

1 loop, best of 3: 210 ms per loop

我们正在考虑是否应该将这种交互式便利设置为默认值,使其比 SciPy 的默认值快 5 倍而不是 20 倍(尽管我们的默认值通常比默认值准确得多SciPy 使用,但那是另一次的数据,可以在基准测试中找到或直接询问)。一方面,它易于使用是有意义的,但另一方面,如果重新启用长时间计算的专业化并且 Monte Carlo 未知(这是你真正想要速度的地方),那么很多那里的人会受到 2 到 4 倍的性能损失,这可能需要额外的 days/weeks 计算。嗯...艰难的选择。

所以最后混合了优化选择和 Julia 缺少的一些预编译功能,这些功能影响了交互性而不影响真正的 运行 时间速度。如果您希望使用一些大的 Monte Carlo 来估计参数,或者求解大量的 SDE,或者求解一个大的 PDE,我们都有。这是我们的第一个目标,我们确保尽可能好地实现它。但是在 REPL 中玩耍确实有 2-3 秒的“故障”,我们也不能忽视(当然比在 C/Fortran 中玩耍更好,但对于 REPL 来说仍然不理想)。为此,我已经向您展示了已经开发和测试的解决方案,因此希望明年这个时候我们可以针对该特定案例提供更好的答案。

PS

还有两件事要注意。如果您只使用 ODE 求解器,则只需执行 using OrdinaryDiffEq 即可保留 downloading/installing/compiling/importing 所有 DifferentialEquations.jl (this is described in the manual)。此外,像这样使用 saveat 可能不是解决此问题的最快方法:在这里用更少的点解决它并根据需要使用密集输出可能更好。

编辑

I opened an issue detailing how we can reduce the "between function" compilation time without losing the speedup that specializing gives。我认为这是我们可以作为短期优先事项的事情,因为我同意我们可以在这里做得更好。