Gumbel dlshogiを作るその２（訓練データ出力処理）

前回の記事で、Gumbel AlphaZeroの自己対局処理の実装を完了した。今回は、その続きとして、訓練データ出力処理を実装した。

訓練データの形式

AlphaZeroでは探索した手のみの訪問回数に応じた分布を学習する。一方、Gumbel AlphaZeroでは、改善されたポリシーには合法手すべてが含まれるため、すべての手の確率を保存する必要がある。

合法手は局面により数が異なるが、Pythonで実装する方針のため、可変長のデータ形式にすると、合法手の数分の繰り返し処理が必要になり訓練の速度に影響する。そこで、非合法手も含めたポリシーをそのまま固定長で保存して、読み込んだ値をそのまま学習できるようにする。

また、固定長にすることで、メモリに展開しないでファイルのランダムアクセスが可能になり、データローダの並列化が行いやすくなる。

非合法手も含めて保存することで、データ量が大きくなるが、dlshogiのように過去分のデータで教師ありを学習は行わない予定であり、古いデータは破棄するのでデータ量は気にしないことにする。

処理方式

各アクターが並列で対局を進行し、終局した対局のデータをファイルに出力する。アクターがファイル出力を行うと、他のアクターを止めることになるため、データをキューに追加し、ファイルへの書き出しは専用スレッドで処理する。

PythonはGILがあるため、CPUバウンドの処理は並列化されないが、ファイルへの書き出しはほとんどIO処理のため、別スレッドで行うことで効率化される。

実装内容

訓練データの構造定義

まず、訓練データを格納するための構造を定義した。以下の2つの要素から構成される：

TrainingDataクラス: ゲーム終了時に生成される訓練データを格納するデータクラス
dtypeTrainingData: NumPyファイルに保存する際のデータ型定義

@dataclass
class TrainingData:
    hcps: np.ndarray                    # 各局面のHuffman符号化された盤面
    policy_outputs: list[base.PolicyOutput]  # 各局面でのMCTSの政策出力
    turn: int                          # ゲーム終了時の手番
    is_game_over: bool                 # 詰み等での終局か
    is_nyugyoku: bool                  # 入玉宣言での終局か
    is_draw: int                       # 引き分けの種類
    is_max_moves: bool                 # 最大手数到達での終局か

dtypeTrainingData = np.dtype([
    ("hcp", dtypeHcp),                 # 局面データ
    ("result", np.uint8),              # 勝敗結果
    ("policy", np.dtype((np.float32, MOVE_LABELS_NUM))),  # 政策ベクトル
])

Actorクラスでの訓練データ収集

各アクターは対局中に以下の情報を収集する：

hcps: 各局面のHuffman符号化データを格納する配列
policy_outputs: 各局面でのMCTSによる政策出力（行動確率分布）

ゲーム終了時に、これらの情報をTrainingDataオブジェクトとしてキューに送信する。

勝敗判定ロジック

訓練データ出力時に、ゲーム終了理由に応じて勝敗を判定する：

# 詰み or 千日手負け → 相手の勝ち
if training_data.is_game_over or training_data.is_draw == REPETITION_LOSE:
    result = WHITE_WIN if training_data.turn == BLACK else BLACK_WIN

# 入玉宣言 or 千日手勝ち → 手番側の勝ち  
elif training_data.is_nyugyoku or training_data.is_draw == REPETITION_WIN:
    result = BLACK_WIN if training_data.turn == BLACK else WHITE_WIN

# 千日手引き分け → 引き分け
elif training_data.is_draw == REPETITION_DRAW:
    result = DRAW

バックグラウンド書き込み処理

訓練データの書き込みは、メインの自己対局処理をブロックしないよう、専用のワーカースレッドで行う：

def write_training_data(queue: Queue, output_dir: str):
    # タイムスタンプ付きファイル名で出力
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
    filename = f"{timestamp}.data"
    filepath = os.path.join(output_dir, filename)
    
    while True:
        training_data = queue.get()
        if training_data is None:  # 終了シグナル
            break
            
        # 勝敗判定と訓練データ作成
        # ...
        
        # バイナリファイルに追記
        with open(filepath, "ab") as f:
            data.tofile(f)

自己対局処理との統合

メインの自己対局処理では、以下の流れで訓練データ出力を統合した：

訓練データ書き込み用のキューとワーカースレッドを開始
各アクターにキューを渡して、ゲーム終了時にTrainingDataを送信
KeyboardInterrupt時に終了シグナル（None）をキューに送信してスレッドを正常終了

技術的なポイント

効率的なファイル出力

タイムスタンプ付きのファイル名で、実行ごとに別ファイルに出力
NumPyの構造化配列を使用してバイナリ形式で効率的に保存
ファイルを追記モード（"ab"）で開いて、ゲーム終了ごとに逐次書き込み

スレッド安全性

Queueを使用してメインスレッドとワーカースレッド間でデータを安全にやり取り
ワーカースレッドをdaemonに設定して、メインプロセス終了時の自動終了を保証

メモリ効率

対局中は必要最小限のデータ（HCPとポリシー出力）のみ保持
ゲーム終了時に一括してファイルに書き出してメモリを解放

まとめ

前回実装した自己対局処理に、訓練データ出力機能を追加した。これにより、Gumbel AlphaZeroアルゴリズムによる自己対局で生成された訓練データを、効率的にファイル保存できるようになった。

実装のポイントは以下の通り：

構造化された訓練データ形式の定義
適切な勝敗判定ロジック
バックグラウンドでの非同期ファイル書き込み
メモリ効率とスレッド安全性の両立

次回は、この訓練データを使用してニューラルネットワークモデルの学習処理を実装する予定である。

TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。