Julia:优化简单动力系统的仿真

Question

我正在尝试优化简单动态系统的仿真,其中网络的响应及其参数(权重)根据简单的线性方程演变.模拟需要运行数千万个时间步,但网络规模通常很小.因此,性能受矩阵向量产品的约束,而不是临时数组,绑定检查和其他不太明显的因素.由于我是朱莉娅的新手,我很欣赏任何进一步优化性能的提示.

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)

网络规模越小,提升越大.可以在此处找到更新的模拟代码的要点.

更新2

大部分存储器分配是由矩阵矢量产物引起的.因此,为了摆脱那些我通过就地BLAS操作BLAS.genv!取代那些产品,它将时间再减少25%,内存分配减少90%,

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

这里更新了代码.

更新3

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

blas_set_num_threads(1)

网络规模为20的时间增加15%,对于规模为50的网络增益为50%.没有更多的内存分配,新的时间是

1.638287 seconds (11 allocations: 1.266 KB)

同样,可以在此处找到更新的代码.

更新4

我写了一个基本的Cython脚本来比较到目前为止的结果.主要区别在于我不使用任何BLAS调用但是有循环:注入低级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 包肯定是一个很大的贡献,但是为了看到它为你做的肮脏的工作,你可以做这样的事情(来自自述包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.]