【Julia】Mmap.mmapで「配列の配列」を作った話

はじめに

Juliaでバイナリデータを読み込む手順がよくわからなかったので、色々と調べた結果をメモする。
特にメモリに乗らない大きなデータを処理したいので、Mmap.mmapが使えないか試してみた。
環境は以下の通り。

• macOS Mojave 10.14.6
• JupyterLab 1.1.4
• Julia 1.3.1

※下記の内容は随所でテキトーなことをいっているので、情報源としては信頼しないようご注意ください。

目的

やりたいことは大規模な時系列データの読み込みであり、以下のような状況を想定している。

• 複数の連番ファイルに各時刻のデータが1次元配列としてバイナリで保存されている
• データをベクトルや行列として読み込み、色々な線形計算を行いたい
• ただし、データ量が大きく、いくつかのファイルをメモリに読み込むことはできるが、すべてのファイルを1つの行列にまとめるとメモリ不足になる
• 各データは密なベクトルであり、構成される行列も密行列となるため、疎行列用のライブラリは適さない

もちろん、データを逐一読み込みながら処理すれば一応メモリ不足は回避できるが、それだとあからさまに遅そうなのでもう少しスマートにやりたい。そこで、ファイルをメモリにマッピングして利用可能なMmap.mmapで処理できないか試してみる。

テストコード

問題設定

$n \gg m$の縦長の乱数行列$\mathbf{A} \in \mathbb{R}^{n \times m}$を生成し、メモリに保持可能な小さな行列になるよう$\mathbf{A}^{\top} \mathbf{A} \in \mathbb{R}^{m \times m}$を計算してみる。
ここでは、$n = 2,000,000$、$m = 100$とし、数値は正規乱数から生成した。倍精度実数を使う場合、トータルのデータ量は約1.6GBになる。

以下の3つの場合で計算時間とメモリ量を見てみる。

1. mmapを使わず、普通の行列として計算
2. 行列全体を1つのファイルに保存し、mmapで読み込む
3. 行列を列ごとに個別のファイルに保存し、mmapで読み込む

計算時間とメモリ量はBenchmarkTools.jlの@benchmarkマクロを使用して確認した。
なお、行列の生成やファイルの書き込み部分は計測の対象外としている。

データの準備

まずは下記のコードでテストデータを作成する。

using Random

# データ行列
n = 2_000_000
m = 100
A = randn(Float64,(n,m))

# 行列を1つのファイルに書き込む
# ヘッダーとして行数と列数を付加
open("test_mat.bin", "w") do io
    write(io, n)
    write(io, m)
    write(io, A)
end

# 行列を列ごとに個別のファイルに書き込む
# ヘッダーとして行数を付加
for i in 1:m
    open("test_mat$i.bin", "w") do io
        write(io, n)
        write(io, A[:,i])
    end
end

これにより、$n \times m$の行列が保存された1つのファイルと、$n$行のデータベクトルが保存された$m$個のファイルが作成される。

mmapを使わず、普通の行列として計算

まずは普通に計算した場合。行列$\mathbf{A}$はすべてメモリに保持されている。

using LinearAlgebra

function test_mat(A)
    A'*A
    nothing
end

計測は以下のようにするだけ。

@benchmark test_mat(A)

結果は以下の通り。

BenchmarkTools.Trial: 
  memory estimate:  78.20 KiB
  allocs estimate:  2
  --------------
  minimum time:     756.787 ms (0.00% GC)
  median time:      759.918 ms (0.00% GC)
  mean time:        762.470 ms (0.00% GC)
  maximum time:     767.908 ms (0.00% GC)
  --------------
  samples:          7
  evals/sample:     1

当然だが、メモリ割り当てはほとんど生じておらず、後述の結果と比べても計算時間は一番小さかった。ただし、関数の外側で行列をメモリに入れているので、Juliaのプロセスを見ると1.8GB程度のメモリ消費になっていた。そのため、データ量が増えた場合にはやはりこの手法は使えない（実際には数十GB以上のデータを処理したいので）。

行列全体を1つのファイルに保存し、mmapで読み込む

次に行列を保存したファイルからデータを読み込んだ場合を見てみる。

function test_mmap_mat()
    f = open("test_mat.bin", "r")
    row = read(f, Int)
    col = read(f, Int)
    A = Mmap.mmap(f, Matrix{Float64},(row,col))
    close(f)
    A'*A
    nothing
end

結果は以下の通り。

BenchmarkTools.Trial: 
  memory estimate:  79.09 KiB
  allocs estimate:  19
  --------------
  minimum time:     1.146 s (0.00% GC)
  median time:      1.190 s (0.00% GC)
  mean time:        1.183 s (0.00% GC)
  maximum time:     1.214 s (0.00% GC)
  --------------
  samples:          5
  evals/sample:     1

メモリに保持した場合に比べて1.5倍程度の時間がかかっているが、それ以外は同様のように見える。実際のメモリ消費はこちらの方が圧倒的に小さいため、この程度の速度低下は許容範囲ではないだろうか。計算時間の増加割合はデータ量に応じて大きくなるかもしれないが、数倍までならメモリ不足にならない方がいいと思う。

行列を列ごとに個別のファイルに保存し、mmapで読み込む

最後に本命の分割されたファイルを読み込む場合。

function test_mmap_vec(m)
    A = []
    for i in 1:m
        f = open("test_mat$i.bin", "r")
        row = read(f, Int)
        vec = Mmap.mmap(f, Vector{Float64},row)
        push!(A,vec)
        close(f)
    end
    
    # 通常の行列ではないため、愚直に計算
    B = zeros(m,m)
    for j in 1:m
        for i in 1:m
            B[i,j] = dot(A[i],A[j])
        end
    end
    nothing
end

上記のコードでは、ファイルから読み出したベクトルを1つにまとめるためにpush!を使っている。このとき、AのタイプはArray{Float64,2}ではなく、Array{Any,1}となり、その要素としてArray{Float64,1}が格納されている形になる。扱いに注意が必要な形ではあるが、これであればマッピングそのものはかなり高速に実行できる。¹

ちなみにhcatやappend!を使うとあからさまに遅く、メモリ消費もかなり大きかった。これはおそらくベクトルを追加する度にメモリを取り直しているためだと思われる。結果は行列になるので扱いは容易だが、mmapの利点が失われるのでやる意味は薄そうだ。

BenchmarkTools.Trial: 
  memory estimate:  507.03 KiB
  allocs estimate:  22210
  --------------
  minimum time:     8.489 s (0.00% GC)
  median time:      8.489 s (0.00% GC)
  mean time:        8.489 s (0.00% GC)
  maximum time:     8.489 s (0.00% GC)
  --------------
  samples:          1
  evals/sample:     1

行列積の計算は他と比べてかなり遅かった。これは結果が行列になっておらず、ベタ書きのアルゴリズムを使っているせいだと考えられる。ちゃんとした行列積のアルゴリズムを使えば1つのファイルから読み込んだ場合と同程度になるのではないだろうか。

そこで、列ごとに処理できるよう、列ベクトルのノルムを計算したときの時間を見てみる。

function test_mmap_vec2(n,m)
    A = []
    for i in 1:m
        f = open("test_mat$i.bin", "r")
        row = read(f, Int)
        vec = Mmap.mmap(f, Vector{Float64},row)
        push!(A,vec)
        close(f)
    end
    
    # 1列ごとにアクセスする計算に変更
    norm.(A)
    nothing
end

その結果、下記のようにかなり高速化された。

BenchmarkTools.Trial: 
  memory estimate:  118.86 KiB
  allocs estimate:  2313
  --------------
  minimum time:     522.845 ms (0.00% GC)
  median time:      641.839 ms (0.00% GC)
  mean time:        641.067 ms (0.00% GC)
  maximum time:     765.137 ms (0.00% GC)
  --------------
  samples:          8
  evals/sample:     1

そのため、既存のライブラリが有効利用できるベクトル単位での処理なら問題なく使えるのだと思われる。とはいえ、汎用的ではないため、扱いやすい小規模な行列に落とし込むための工夫が必要だろう。この辺は今後の課題だ。

おわりに

あんまり良さげな結果でもないが、mmapで分割されたファイルを読み込むことはできたので、ひとまずはこれでよしとしよう。
push!で作った「配列の配列」は正直使い勝手が悪いので、素直に中間ファイルに書き出してから改めて行列としてマッピングするのもありなんじゃないかと思ったが、数十GB単位のファイルコピーなんてさすがにやりたくない。
なので、列ベクトルだけが参照できる状態でも計算可能な低ランク近似アルゴリズムを適用し、早々に生データにアクセスしないで済むようにすべきだろう。これはこれで興味があるので、また色々調べてやってみようと思う。

1. 後で確認したらA = Vector{Float64}[]として初期化すると結果はArray{Array{Float64,1},1}になるので、型を指定した方が分かりやすいとは思う。 ↩