SQLiteによる教師データの管理

先日AlphaZero方式で教師データを生成する際に、データを固定サイズにすることを検討した。
しかし、指し手の確率分布を保存するには、合法手500手近くの領域が必要となるため、1回の訓練ステップ全てのデータをメモリに載せるのは厳しいことがわかった。
AlphaZeroでは、過去ゲーム数100万からランダムサンプリングするので、1ゲーム平均80手とすると、訪問数を2byteで確保すると、500手×80手×100万×2byte=80GBが必要になる。

現在実験しているdlshogiでは、教師データは単一の指し手のみ保存しているため、一度に500万局面を学習しても必要メモリは2GBに収まる。

そこで、すべてメモリに読み込む方式はあきらめ、ファイルからランダムアクセスを行うようにする。
ゲームごとに局面数が異なるため、ゲーム数をウィンドウサイズとした場合、ゲームごとの局面数を管理しておく必要がある。
局面のデータとゲームごとの局面数の管理を自前で行うのも面倒なので、データベースを使用することにした。
データベースにサーバが必要になると煩雑になるのでPythonで標準で使えるSQLiteで実装することにした。
将来的にスケールアウトしたい場合もSQLiteから他のRDBMSに移行するのは難しくない。

ランダムサンプリング

SQLiteで、ゲーム数をウィンドウサイズとしてランダムサンプリングするには、単純なSQLでは、ORDER BY RANDOM() LIMIT Nで行えるが、これはデータ数が多い場合に非常に遅くなる。

WITH RECURSIVEを使う

SQLiteは、WITH RECURSIVE構文で仮想なテーブルを再帰的に定義できる。
The WITH Clause
これを使うと、関数型プログラミングのようなことができる。

ランダムで100個の数値を得るには以下のように行う。

WITH RECURSIVE
    rand_id(id) AS (
        SELECT RANDOM()
        UNION ALL SELECT RANDOM() FROM rand_id
        LIMIT 100)
SELECT ABS(id) FROM rand_id

これを使うと、ウィンドウサイズを指定したランダムサンプリングが実現できる。

WITH RECURSIVE
    rand_id(id) AS (
        SELECT RANDOM()
        UNION ALL SELECT RANDOM() FROM rand_id
        LIMIT {n_samples})
SELECT hcps, repetition, total_move_count, legal_moves, visits, game_result FROM training_data AS A
    JOIN (
        SELECT ABS(id) % ((SELECT MAX(rowid) FROM training_data) - (SELECT MIN(rowid) FROM training_data WHERE game_id >= {min_game_id})) + (SELECT MIN(rowid) FROM training_data WHERE game_id >= {min_game_id}) AS B
        FROM rand_id)
    ON A.rowid = B

ただし、rowidが連続していることを前提としている。重複も許容している。
重複を排除したい場合はDISTINCTをうまく使えばできる。

SQLiteを使う場合は、固定サイズにこだわる必要はないため、合法手の分だけ保存すればよいのでディスクサイズの節約にもなる。