提高 Julia 中自定义函数性能的技巧

Question

我正在使用 Julia 复制最初在 Matlab 中制作的一系列步骤。在 Octave 中，此过程需要 1.4582 秒，而在 Julia（使用 Jupyter）中，大约需要 10 秒。我会尽量在脚本中保持简短。我的目标是达到或提高 Octave 的性能。首先，我将描述我的变量和一些函数：

1. zgrid（双 1x7 尺寸）
2. kgrid（双倍 500x1 尺寸）
3. V0（双倍 500x7 尺寸）
4. P（双 7x7 大小）一个转换矩阵
5. delta 和 beta 是固定参数。
6. F(z,k) 和 u(c) 是特殊函数，在 Julia 脚本中指定。

% Octave script
% V0 is given
[K, Z, K2] = meshgrid(kgrid, zgrid, kgrid);
K = permute(K, [2, 1, 3]);
Z = permute(Z, [2, 1, 3]);
K2 = permute(K2, [2, 1, 3]);

C = max(f(Z,K) + (1-delta)*K - K2,0);
U = u(C);

EV = V0*P';% EV is a 500x7 matrix size
EV = permute(repmat(EV, 1, 1, 500), [3, 2, 1]);
H = U + beta*EV;
[TV, index] = max(H, [], 3);

在 Julia 中，我创建了一个复制此过程的函数。我使用了循环，但它的性能要长 9 倍。

% Julia script
% V0 is the input of my T operator function
V0 = repeat(sqrt.(kgrid), outer = [1,7]);
F = (z,k) ->  exp(z)*(k^α);
u = (c) -> (c^(1-μ) - 1)/(1-μ)

% parameters
α = 1/3
β = 0.987
δ = 0.012;
μ = 2
Kss = 48.1905148382166
kgrid = range(0.75*Kss, stop=1.25*Kss, length=500);
zgrid = [-0.06725382459813659, -0.044835883065424395, -0.0224179415327122, 0 , 0.022417941532712187, 0.04483588306542438, 0.06725382459813657]

function T(V)
    E=V*P'
    T1 = zeros(Float64, 500, 7 )
    aux = zeros(Float64, 500)
    for i = 1:7
        for j = 1:500
            for l = 1:500
                c= maximum( (F(zrid[i],kgrid[j]) +(1-δ)*kgrid[j] - kgrid[l],0))
                aux[l] =  u(c) + β*E[l,i]
            end
            T1[j,i] = maximum(aux)     
        end
    end
    return T1
end

我非常想提高我在 Julia 中的表现。我相信有办法改进，但我是 Julia 编程新手。

Answer 1

此代码在 5 毫秒内为我运行。请注意，我已将 F 和 u 设置为适当的（非匿名）函数 F_ 和 u_，但您可以通过创建匿名函数 const 来获得类似的效果。

您的主要问题是您有很多非const 全局变量，并且您的主函数多次执行不必要的工作，并创建了一个不必要的数组aux。

手册中的性能提示部分是必不可少的阅读：https://docs.julialang.org/en/v1/manual/performance-tips/

F_(z,k) = exp(z) * (k^(1/3));  # you can still use α, but it must be const
u_(c) = (c^(1-2) - 1)/(1-2)

function T_(V, P, kgrid, zgrid, β, δ)
    E = V * P'
    T1 = similar(V)
    for i in axes(T1, 2)
        for j in axes(T1, 1)
            temp = F_(zgrid[i], kgrid[j]) + (1-δ)*kgrid[j]
            aux = -Inf
            for l in eachindex(kgrid)
                c = max(0.0, temp - kgrid[l])
                aux = max(aux, u_(c) + β * E[l, i])
            end
            T1[j,i] = aux
        end
    end
    return T1
end

基准测试：

V0 = repeat(sqrt.(kgrid), outer = [1,7]);
zgrid = sort!(rand(1, 7); dims=2)
kgrid = sort!(rand(500, 1); dims=1)
P = rand(length(zgrid), length(zgrid))
@btime T_($V0, $P, $kgrid, $zgrid, $β, $δ);
# output:   5.126 ms (4 allocations: 54.91 KiB)

Answer 2

以下应该表现得更好。最明显的区别在于它计算的F 减少了 500 倍，并且不依赖于全局变量。

function T(V,kgrid,zgrid,β,δ)
    E=V*P'
    T1 = zeros(Float64, 500, 7)
    for j = 1:500
        for i = 1:7
            x = F(zrid[i],kgrid[j]) +(1-δ)*kgrid[j] 
            T1[j,i] = maximum(u(max(x - kgrid[l], 0)) + β*E[l,i] for l in 1:500)
        end
    end
    return T1
end