performance - Julia:优化简单动力系统的模拟

Question

我正在尝试优化一个简单动态系统的模拟，其中网络的响应及其参数(权重)根据简单的线性方程发展。模拟需要运行数千万个时间步，但网络规模通常很小。因此，性能受矩阵向量乘积的限制较少，而是受临时数组、边界检查和其他不太明显的因素的限制。由于我是 Julia 的新手，我会很感激任何进一步优化性能的提示。

function train_network(A, T, Of, cs, dt)
    N, I = size(T)
    z    = zeros(I)
    r    = zeros(N)

    @inbounds for t in 1:size(cs, 1)
        # precompute
        Az  = A*z
        Ofr = Of*r

        # compute training signal
        @devec z += dt.*(Az + cs[t] - 0.5.*z)
        I_teach   = T*(Az + cs[t])
        Tz        = T*z

        # rate updates
        @devec r += dt.*(I_teach - Ofr - 0.1.*r)

        # weight updates
        for i in 1:I
            @devec T[:, i] += dt.*1e-3.*(z[i].*r - T[:, i])
        end

        for n in 1:N
            @devec Of[:, n] += dt.*1e-3.*(Tz.*r[n] - Of[:, n])     
        end
    end
end

# init parameters
N, I = 20, 2
dt  = 1e-3

# init weights
T = rand(N, I)*N
A = rand(I, I)
Of = rand(N, N)/N

# simulation time & input
sim_T = 2000
ts = 0:dt:sim_T
cs = randn(size(ts, 1), I)

定时网络(2.000.000 步)

@time train_network(A, T, Of, cs, dt)

产生时间

3.420486 seconds (26.12 M allocations: 2.299 GB, 6.65% gc time)

更新 1

按照 David Sanders 的建议，我摆脱了 devec 宏并写出了循环。这确实减少了阵列分配并将性能提高了约 25%，以下是新数字:

2.648113 seconds (18.00 M allocations: 1.669 GB, 5.60% gc time)

网络规模越小，提升越大。可以找到更新的模拟代码的要点 here .

更新 2

大部分内存分配是由于矩阵向量乘积造成的。因此，为了摆脱那些我用就地 BLAS 操作 BLAS.genv! 替换了这些产品，这将时间再减少了 25%，内存分配减少了 90%，

1.990031 seconds (2.00 M allocations: 152.589 MB, 0.69% gc time)

更新代码 here .

更新 3

最大的 rank-1 更新也可以被两个就地 BLAS 函数调用替换，即 BLAS.scal!用于缩放和 BLAS.ger!进行 rank-1 更新。需要注意的是，如果使用多个线程(OpenBLAS 有问题？)，这两个调用都相当慢，所以最好设置

blas_set_num_threads(1)

网络规模为 20 的时间增加 15%，网络规模为 50 的时间增加 50%。不再分配内存，新的时间为

1.638287 seconds (11 allocations: 1.266 KB)

同样，更新后的代码可以在 here 找到.

更新 4

我写了一个基本的 Cython script比较迄今为止的结果。主要区别在于我不使用任何对 BLAS 的调用，但有循环:注入(inject)低级 BLAS 调用在 Cython 中很痛苦，而对 numpy dot 的调用对于小型网络来说开销太大(我试过......)。时间是

CPU times: user 3.46 s, sys: 6 ms, total: 3.47 s, Wall time: 3.47 s

这与原始版本大致相同(到目前为止，已削减了 50%)。

Answer 1

虽然您使用的是 Devectorize.jl包，我建议您将所有这些矢量化操作显式写为简单的循环。我希望这会给您带来显着的性能提升。
Devectorize package 确实是一个很大的贡献，但是要看到它为你做的脏活，你可以做这样的事情(来自 package README 的一个例子):

using Devectorize

a = rand(2,2);
b = rand(2,2);
c = rand(2,2);

julia> macroexpand(:(@devec r = exp(a + b) .* sum(c)))

在这里， macroexpand是一个函数，它告诉您 @devec 的代码。宏扩展其参数(该行其余部分的代码)。
我不会在这里显示输出来破坏惊喜，但它不仅仅是简单的 for你会用手写的循环。

此外，您拥有大量分配的事实表明并非所有向量操作都得到正确处理。

顺便说一句，不要忘记先做一个小运行，这样你就不会为编译步骤计时。

[切线注释:这里， exp是将通常的指数函数应用于矩阵的每个元素的函数，相当于 map(exp, a+b) . expm给出矩阵的指数.有人说不赞成使用 exp .]