TadaoYamaokaの日記

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

cuDNNでAlphaGoのSL policy networkの順伝播を実装する

Chainerを使って学習したAlphaGoのSL policy networkのモデルを使用して、C++のプログラムからcuDNNを使用して順伝播を実装してみた。

SL policy networkを囲碁のプログラムに組み込んで使おうとすると、PythonとChainerが必要になるのでは実行環境の敷居が高くなるので、できればC++のみで実装したい。
実行環境では、学習を行う必要がないので、順伝播のみを実行できればよい。
順伝播のみであれば、cuDNNを直接利用して比較的簡単に実装できる。

AlphGoのSL policiy networkは、cuDNNを使用すると以下のように実装できる。

	// 順伝播計算
	// layer1
	// 畳み込み
	const float alpha = 1.0f;
	const float beta = 0.0f;
	checkCUDNN(cudnnConvolutionForward(cudnnHandle,
		&alpha,
		srcTensorDesc,
		srcData_dev,
		filterDesc1,
		filterData_dev1,
		convDesc1,
		algo1,
		workSpace1,
		workSpaceSizeInBytes1,
		&beta,
		hdnTensorDesc,
		hdnData_dev1));
	// バイアス
	const float alpha_bias = 1.0f;
	const float beta_bias = 1.0f;
	checkCUDNN(cudnnAddTensor(cudnnHandle, &alpha_bias, biasTensorDesc, biasData_dev1, &beta_bias, hdnTensorDesc, hdnData_dev1));
	// 活性化関数
	checkCUDNN(cudnnActivationForward(cudnnHandle,
		activDesc,
		&alpha,
		hdnTensorDesc,
		hdnData_dev1,
		&beta,
		hdnTensorDesc,
		hdnData_dev1));

	// layer2
	// 畳み込み
	checkCUDNN(cudnnConvolutionForward(cudnnHandle,
		&alpha,
		hdnTensorDesc,
		hdnData_dev1,
		filterDesc2,
		filterData_dev2,
		convDesc2,
		algo2,
		workSpace2,
		workSpaceSizeInBytes2,
		&beta,
		hdnTensorDesc,
		hdnData_dev2));
	// バイアス
	checkCUDNN(cudnnAddTensor(cudnnHandle, &alpha_bias, biasTensorDesc, biasData_dev2, &beta_bias, hdnTensorDesc, hdnData_dev2));
	// 活性化関数
	checkCUDNN(cudnnActivationForward(cudnnHandle,
		activDesc,
		&alpha,
		hdnTensorDesc,
		hdnData_dev2,
		&beta,
		hdnTensorDesc,
		hdnData_dev2));

(略)

	// layer13
	// 畳み込み
	cudnnConvolutionForward(cudnnHandle,
		&alpha,
		hdnTensorDesc,
		hdnData_dev2,
		filterDesc13,
		filterData_dev13,
		convDesc13,
		algo13,
		workSpace13,
		workSpaceSizeInBytes13,
		&beta,
		dstTensorDesc,
		hdnData_dev13);
	// バイアス
	checkCUDNN(cudnnAddTensor(cudnnHandle, &alpha_bias, biasTensorDesc13, biasData_dev13, &beta_bias, dstTensorDesc, hdnData_dev13));
	// 活性化関数
	cudnnSoftmaxForward(cudnnHandle,
		CUDNN_SOFTMAX_FAST,
		CUDNN_SOFTMAX_MODE_INSTANCE,
		&alpha,
		dstTensorDesc,
		hdnData_dev13,
		&beta,
		dstTensorDesc,
		dstData_dev);

上記コードは初期化処理などは省いている。全ソースはGithubを参照してほしい。

layer2~layer12は同じ畳み込みとReLUの繰り返しになる。
GPUのメモリを節約するために、中間層の出力メモリは同じ領域を使いまわしている。
畳み込みは入力のメモリと出力のメモリは同じ領域にできないので、layerごとに交互に違う領域を使用している。

Chainerで学習したモデルの読み込みは、以前の日記で書いた方法で実装した。


特徴をAlphaGo論文のExtended Data Table 2の上の4つで学習したChainerのモデルを使用して、Chainerで実装した順伝播と同じ入力を与えて比較したところ、同じ出力になることを確認した。


これを囲碁プログラムのツリーポリシーに組み込んで、どれだけ効果があるか確認したい。