TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

NNの計算結果をキャッシュする

自己対局でノードの再利用しないようにしたが、NNの計算結果は再利用した方が効率がよいため、キャッシュの仕組みを導入したい。

並列で実行しているゲームすべてについて1ゲーム分のNNの結果を保持するにはメモリ容量が不足するため、使用されなくなった局面(古い手番で探索した局面)を削除する仕組みが必要になる。
その仕組みとして、単純な方法としては、ゲームごとに開始局面からの手数を記録しておき、局面の探索開始前に古い手数の局面を削除するようなことが考えられる。
しかし、一部のエントリを削除する場合、エントリの線形探索やハッシュエントリの再配置が必要になるため効率が良くない。

LRUキャッシュ

古いエントリを効率よく削除するデータ構造として、Least Recently Used(LRU) Cacheという仕組みがある。
オペレーティングシステムの教科書とかでよく出てくる。

C++のLRU Cacheの実装には、検索すればサンプルコードがいくつも出てくる。
ほとんどは、mapかunordered_mapと、リンクリストを使って実装されている。
boostでもlru_cache.hppに実装されている。
ただし、並列実行する自己対局で使用するには、スレッドセーフな実装が必要になる。


Leela Chess Zeroでも、NNのキャッシュの仕組みとしてLeast Recently Used(LRU) Cacheが使用されている。
スレッドセーフで、ゾブリストハッシュと組み合わせて使用できるように実装されている。
テンプレートで汎用的に実装されているため、将棋AIでもそのまま流用できそうなので、利用させてもらうことにする。

Leela Chess ZeroのLRU Cacheは以下のようにして使用する。

#include <iostream>
#include <vector>
#include <random>
#include "LruCache.h"

constexpr int cache_size = 8;

int main() {
	LruCache<int, float> cache(cache_size);

	std::mt19937 rnd;
	std::vector<int> key(10);
	for (int i = 0; i < 10; i++)
		key[i] = rnd();

	for (int i = 0; i < 10; i++)
		cache.Insert(key[i], std::move(std::make_unique<float>(float(i))));

	for (int i = 0; i < 10; i++) {
		LruCacheLock<int, float> cache_lock(&cache, key[i]);
		auto v = *cache_lock;
		std::cout << i << "\t";
		if (v == nullptr)
			std::cout << "null" << std::endl;
		else
			std::cout << *v << std::endl;
	}
}

※LruCache.hが、Leela Chess Zeroのcache.hを流用したソース

LRUキャッシュを作成するには、LruCacheのテンプレート引数の1つ目にキーの型、2つ目に格納する値の型を指定する。コンストラクタの引数にはキャッシュサイズを指定する。

Insertメソッドにキーと値を設定して、エントリを追加する。
値は、unique_ptrにする必要がある。

値を検索して使用する場合は、LookupAndPinメソッドで値を取り出して使用し、使用が終わった後Unpinで値を戻す。
LruCacheLockを使うとスコープ範囲内で、LookupAndPinとUnpinの呼び出しを自動化できる。

実行結果

上記コードの実行結果は以下のようになる。

0       null
1       null
2       2
3       3
4       4
5       5
6       6
7       7
8       8
9       9

キャッシュサイズが8で、10エントリを追加しており、古いエントリから削除されていることが確認できる。