提高 Julia 中 for 循环的速度

Question

这是我在 Julia 平台中的代码，我想加快速度。无论如何，我可以让这个更快吗？50k*50k的数据集需要0.5秒。我原以为 Julia 会比这快得多，否则我不确定我是否在做一个愚蠢的实现。

ar = [[1,2,3,4,5], [2,3,4,5,6,7,8], [4,7,8,9], [9,10], [2,3,4,5]]

SV = rand(10,5)

function h_score_0(ar ,SV)
    m = length(ar)
    SC = Array{Float64,2}(undef, size(SV, 2), m)
    for iter = 1:m
        nodes = ar[iter]
        for jj = 1:size(SV, 2)
            mx = maximum(SV[nodes, jj])
            mn = minimum(SV[nodes, jj])
            term1 = (mx - mn)^2;
            SC[jj, iter] = (term1);
        end
    end
    return score = sum(SC, dims = 1)
end

Answer 1

您的代码中有一些不必要的分配：

\n

mx = maximum(SV[nodes, jj])\nmn = minimum(SV[nodes, jj])\n

\n

切片分配，因此每行都会在此处复制数据，您实际上将数据复制两次，每行一次。您可以确保仅复制一次，甚至更好：使用view，这样就根本没有复制（请注意，view如果您使用的是旧版本，则在 Julia v1.5 上速度要快得多）。

\n

SC = Array{Float64,2}(undef, size(SV, 2), m)\n

\n

没有理由在这里创建一个矩阵，然后对其求和，只需在迭代时进行累加：

\n

score[i] += (mx - mn)^2\n

\n

对于您指定的输入数据，这里的函数在我的笔记本电脑上速度快 5 倍：

\n

function h_score_1(ar, SV)\n    score = zeros(eltype(SV), length(ar))\n    @inbounds for i in eachindex(ar)\n        nodes = ar[i]\n        for j in axes(SV, 2)\n            SVview = view(SV, nodes, j)\n            mx = maximum(SVview)\n            mn = minimum(SVview)\n            score[i] += (mx - mn)^2\n        end\n    end\n    return score\nend\n

\n

该函数输出一个一维向量，而不是原始函数中的 1xN 矩阵。

\n

原则上，如果我们替换，这可能会更快

\n

mx = maximum(SVview)\nmn = minimum(SVview)\n

\n

和

\n

(mn, mx) = extrema(SVview)\n

\n

它只遍历向量一次，而不是两次。不幸的是，存在性能问题extrema，因此它目前不如单独maximum/minimum调用那么快： https: //github.com/JuliaLang/julia/issues/31442

\n

最后，为了绝对以简洁为代价获得最佳性能，我们可以完全避免创建视图，并将调用转变maximum为minimum单个显式循环遍历：

\n

function h_score_2(ar, SV)\n    score = zeros(eltype(SV), length(ar))\n    @inbounds for i in eachindex(ar)\n        nodes = ar[i]\n        for j in axes(SV, 2)\n            mx, mn = -Inf, +Inf\n            for node in nodes\n                x = SV[node, j]\n                mx = ifelse(x > mx, x, mx)\n                mn = ifelse(x < mn, x, mn)\n            end\n            score[i] += (mx - mn)^2\n        end\n    end\n    return score\nend\n

\n

这也避免了所遭受的性能问题，并且每个节点extrema查找一次元素。SV尽管这个版本写起来很烦人，但它的速度要快得多，即使在视图免费的 Julia 1.5 上也是如此。以下是您的测试数据的一些基准计时：

\n

julia> using BenchmarkTools\n\njulia> @btime h_score_0($ar, $SV)\n  2.344 \xce\xbcs (52 allocations: 6.19 KiB)\n1\xc3\x975 Matrix{Float64}:\n 1.95458  2.94592  2.79438  0.709745  1.85877\n\njulia> @btime h_score_1($ar, $SV)\n  392.035 ns (1 allocation: 128 bytes)\n5-element Vector{Float64}:\n 1.9545848011260765\n 2.9459235098820167\n 2.794383144368953\n 0.7097448590904598\n 1.8587691646610984\n\njulia> @btime h_score_2($ar, $SV)\n  118.243 ns (1 allocation: 128 bytes)\n5-element Vector{Float64}:\n 1.9545848011260765\n 2.9459235098820167\n 2.794383144368953\n 0.7097448590904598\n 1.8587691646610984\n

\n

因此，在这里显式地写出最内层的循环是值得的，可以将时间减少 3 倍左右。令人烦恼的是，Julia 编译器尚无法生成如此高效的代码，但它确实在每个版本中都变得更加智能。另一方面，显式循环版本将永远很快，因此如果这段代码确实对性能至关重要，那么可能值得这样写。

\n