Julia

Juliaのコードに対するちょっとしたコメント

@ksknwさんのJuliaとPythonで擬似コードや数式を動かそうとしたときの比較という記事について,Juliaのコードがあまりパフォーマンスの出る書き方ではなかったので,少しコメントしたいと思います。元記事のコメント欄に書くにはちょっと長すぎる気がするのでここに書きました。なお,元のコードのライセンスが指定されていませんでしたので,コードを転載したり修正したりするのはあまり良くないですが,ブログ記事のちょっとしたコードですので許されると思ってこちらにその一部を転載します。元記を書かれた方は,問題がありましたらご指摘下さい。すぐに削除します。

元のコードはこんな感じです。

function minVal(v₁, v₂, v₃)
    if first(v₁)  minimum([first(v₂), first(v₃)])
        return v₁, 1
    elseif first(v₂)  first(v₃)
        return v₂, 2
    else
        return v₃, 3
    end
end

function DTW(A, B)
    S = length(A)
    T = length(B)
    m = Matrix(S, T)
    m[1, 1] = [δ(A[1], B[1]), (0,0)]
    for i in 2:S
        m[i,1] = [m[i-1, 1][1] + δ(A[i], B[1]), [i-1, 1]]
    end
    for j in 2:T
        m[1,j] = [m[1, j-1][1] + δ(A[1], B[j]), [1, j-1]]
    end
    for i in 2:S
        for j in 2:T
            min, index = minVal(m[i-1,j], m[i,j-1], m[i-1,j-1])
            indexes = [[i-1, j], [i, j-1], [i-1, j-1]]
            m[i,j] = first(min) + δ(A[i],B[j]), indexes[index]
        end
    end
    return m
end

私がちょっと書き直したコードがこんな感じです。

function minVal2(v₁, v₂, v₃)
    if first(v₁)  min(first(v₂), first(v₃))
        return v₁, 1
    elseif first(v₂)  first(v₃)
        return v₂, 2
    else
        return v₃, 3
    end
end

function DTW2(A, B)
    S = length(A)
    T = length(B)
    m = Matrix{Tuple{typeof(δ(A[1], B[1])),Tuple{Int,Int}}}(S, T)
    m[1, 1] = (δ(A[1], B[1]), (0,0))
    for i in 2:S
        m[i,1] = (m[i-1, 1][1] + δ(A[i], B[1]), (i-1, 1))
    end
    for j in 2:T
        m[1,j] = (m[1, j-1][1] + δ(A[1], B[j]), (1, j-1))
    end
    for i in 2:S
        for j in 2:T
            min, index = minVal(m[i-1,j], m[i,j-1], m[i-1,j-1])
            indexes = ((i-1, j), (i, j-1), (i-1, j-1))
            m[i,j] = first(min) + δ(A[i],B[j]), indexes[index]
        end
    end
    return m
end

主な変更点としては,
- Matrixの要素の型が推論できるようにした
- ベクトルの代わりにタプルを使った

簡単にベンチマークしてみると,13倍くらいは速くなりました。
メモリ周りでもうちょっと頑張れるところもありますが,とりあえずは良いでしょう。

julia> using BenchmarkTools

julia> @benchmark DTW(A, B)
BenchmarkTools.Trial:
  memory estimate:  661.99 MiB
  allocs estimate:  14917884
  --------------
  minimum time:     2.038 s (58.12% GC)
  median time:      2.427 s (63.63% GC)
  mean time:        2.554 s (65.93% GC)
  maximum time:     3.196 s (72.65% GC)
  --------------
  samples:          3
  evals/sample:     1

julia> @benchmark DTW2(A, B)
BenchmarkTools.Trial:
  memory estimate:  144.71 MiB
  allocs estimate:  2994005
  --------------
  minimum time:     100.388 ms (13.39% GC)
  median time:      179.648 ms (48.18% GC)
  mean time:        147.746 ms (37.73% GC)
  maximum time:     195.191 ms (50.67% GC)
  --------------
  samples:          34
  evals/sample:     1

2つ目のStochastic Block Modelの方はデータがなくてちょっと再現できないのですが,ここのコードを見る限り,グローバル変数に問題がありそうです。ここでは,

K = 8
α = 6
a₀ = b₀ = 0.5

α₁ = transpose(ones(K) * α) # =α₂
logΓ = lgamma

と書いてあるのですが,定数でないグローバル変数にはJuliaの型推論がはたらかず速いコードを生成できないので,以下のように全てにconstをつけると速くなると思います。

const K = 8
const α = 6
const a₀ = 0.5
const b₀ = 0.5

const α₁ = transpose(ones(K) * α) # =α₂
const logΓ = lgamma

Juliaのコードは,型推論がちゃんと働くように書くことが非常に重要です。これに気をつけるだけで10-100倍のパフォーマンスの違いが出ますので,速いコードを書きたい人は是非その辺を気をつけてみて下さい。