TadaoYamaokaの開発日記

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

Chainerで学習したモデルを使ってcuDNNで推論する(BatchNormalization)

C++ CUDA Chainer cuDNN

前回実装した、Chainerで学習したモデルを使用してcuDNNで推論するコードに、BatchNormalizationを追加した。

BatchNormalizationも、cuDNNにAPIが用意されているため、簡単に使用できる。

ネットワーク定義(Chainer)

まず、Chainerで学習するモデルにBatchNormalizationを追加した。

nn.py
from chainer import Chain
import chainer.functions as F
import chainer.links as L

# ネットワーク定義
k = 16
fcl = 256
class NN(Chain):
    def __init__(self):
        super(NN, self).__init__()
        with self.init_scope():
            self.conv1 = L.Convolution2D(in_channels = 1, out_channels = k, ksize = 3, pad = 1)
            self.conv2 = L.Convolution2D(in_channels = k, out_channels = k, ksize = 3, pad = 1)
            self.l3    = L.Linear(7*7*k, fcl)
            self.l4    = L.Linear(fcl, 10)
            self.bn1   = L.BatchNormalization(k)
            self.bn2   = L.BatchNormalization(k)

    def __call__(self, x):
        h = self.conv1(F.reshape(x, (len(x), 1, 28, 28)))
        h = self.bn1(h)
        h = F.max_pooling_2d(F.relu(h), 2)
        h = self.conv2(h)
        h = self.bn2(h)
        h = F.max_pooling_2d(F.relu(h), 2)
        h = F.relu(self.l3(h))
        return self.l4(h)

BatchNormalization層の実装(C++)

cuDNNを使用して推論を行うC++のコードにBatchNormalizationのレイヤーを追加した。

layers.h
template<const int k>
class BatchNormalization {
public:
	BatchNormalization() : bnScale(nullptr), bnBias(nullptr), estimatedMean(nullptr), estimatedVariance(nullptr) {
		const size_t size = k;
		checkCudaErrors(cudaMalloc((void**)&bnScale, size * sizeof(float)));
		checkCudaErrors(cudaMalloc((void**)&bnBias, size * sizeof(float)));
		checkCudaErrors(cudaMalloc((void**)&estimatedMean, size * sizeof(float)));
		checkCudaErrors(cudaMalloc((void**)&estimatedVariance, size * sizeof(float)));
	}
	~BatchNormalization() {
		checkCudaErrors(cudaFree(bnScale));
		checkCudaErrors(cudaFree(bnBias));
		checkCudaErrors(cudaFree(estimatedMean));
		checkCudaErrors(cudaFree(estimatedVariance));
	}

	void operator() (cudnnHandle_t handle, cudnnTensorDescriptor_t xDesc, float* x, float* y) {
		const float alpha = 1.0f;
		const float beta = 0.0f;
		const double eps = 2e-5;
		checkCUDNN(cudnnDeriveBNTensorDescriptor(bnScaleBiasMeanVarDesc, xDesc, CUDNN_BATCHNORM_SPATIAL));
		checkCUDNN(cudnnBatchNormalizationForwardInference(handle, CUDNN_BATCHNORM_SPATIAL, &alpha, &beta, xDesc, x, xDesc, y, bnScaleBiasMeanVarDesc, bnScale, bnBias, estimatedMean, estimatedVariance, eps));
	}

	void set_param(float* bnScale, float *bnBias, float *estimatedMean, float *estimatedVariance) {
		const size_t size = k;
		checkCudaErrors(cudaMemcpy(this->bnScale, bnScale, size * sizeof(float), cudaMemcpyHostToDevice));
		checkCudaErrors(cudaMemcpy(this->bnBias, bnBias, size * sizeof(float), cudaMemcpyHostToDevice));
		checkCudaErrors(cudaMemcpy(this->estimatedMean, estimatedMean, size * sizeof(float), cudaMemcpyHostToDevice));
		checkCudaErrors(cudaMemcpy(this->estimatedVariance, estimatedVariance, size * sizeof(float), cudaMemcpyHostToDevice));
	}

private:
	CudnnTensorDescriptor bnScaleBiasMeanVarDesc;
	float *bnScale;
	float *bnBias;
	float *estimatedMean;
	float *estimatedVariance;
};

ネットワーク定義(C++)

ネットワークの定義に、上記で実装したBatchNormalizationの層を追加した。

nn.h
#pragma once

#include "layers.h"

const int IMAGE_H = 28;
const int IMAGE_W = 28;

const int batch_size = 2;


class NN {
public:
	typedef float x_t[batch_size][1][IMAGE_H][IMAGE_W];
	typedef float y_t[batch_size][10];

	NN();
	~NN();

	void load_model(const char* filename);

	void foward(x_t x, y_t y);

private:
	static CudnnHandle cudnnHandle;
	static CublasHandle cublasHandle;
	static const int k = 16;
	static const int fcl = 256;

	ConvLayer<k, 1, 3, 1> conv1;
	Bias<k, 1, 1> bias1;
	ConvLayer<k, k, 3, 1> conv2;
	Bias<k, 1, 1> bias2;
	Linear<7 * 7 * k, fcl> l3;
	Bias<fcl, 1, 1> bias3;
	Linear<fcl, 10> l4;
	Bias<10, 1, 1> bias4;
	BatchNormalization<k> bn1;
	BatchNormalization<k> bn2;

	ReLU relu;
	MaxPooling2D<2> max_pooling_2d;

	CudnnTensorDescriptor xDesc;
	CudnnTensorDescriptor h1Desc;
	CudnnTensorDescriptor h2Desc;
	CudnnTensorDescriptor h3Desc;
	CudnnTensorDescriptor h4Desc;
	CudnnTensorDescriptor h5Desc;
	CudnnTensorDescriptor yDesc;

	float* x_dev;
	float* h1_dev;
	float* h1_bn_dev;
	float* h2_dev;
	float* h3_dev;
	float* h3_bn_dev;
	float* h4_dev;
	float* h5_dev;
	float* y_dev;
};

パラメータ読み込み

Chainerで学習したモデルの読み込みにBatchNormalizationのパラメータの読み込みを追加した。

nn.cpp
void NN::load_model(const char* filepath)
{
	// load nn params
	ParamMap params;
	load_npz(filepath, params);

	conv1.set_param(params["conv1/W.npy"].data);
	bias1.set_bias(params["conv1/b.npy"].data);
	conv2.set_param(params["conv2/W.npy"].data);
	bias2.set_bias(params["conv2/b.npy"].data);
	l3.set_param(params["l3/W.npy"].data);
	bias3.set_bias(params["l3/b.npy"].data);
	l4.set_param(params["l4/W.npy"].data);
	bias4.set_bias(params["l4/b.npy"].data);
	bn1.set_param(params["bn1/gamma.npy"].data, params["bn1/beta.npy"].data, params["bn1/avg_mean.npy"].data, params["bn1/avg_var.npy"].data);
	bn2.set_param(params["bn2/gamma.npy"].data, params["bn2/beta.npy"].data, params["bn2/avg_mean.npy"].data, params["bn2/avg_var.npy"].data);
}

Chainerで学習したモデルのBatchNormalization層は、gamma.npy、beta.npy、avg_mean.npy、avg_var.npyというファイル名で格納されている。

推論

推論の処理にBatchNormalization層を追加した。

void NN::foward(x_t x, y_t y)
{
	// input
	checkCudaErrors(cudaMemcpy(x_dev, x, sizeof(x_t), cudaMemcpyHostToDevice));

	// conv1
	conv1(cudnnHandle, xDesc, x_dev, h1Desc, h1_dev);
	bias1(cudnnHandle, h1Desc, h1_dev);
	bn1(cudnnHandle, h1Desc, h1_dev, h1_bn_dev);
	relu(cudnnHandle, h1Desc, h1_bn_dev);
	max_pooling_2d(cudnnHandle, h1Desc, h1_bn_dev, h2Desc, h2_dev);

	// conv2
	conv2(cudnnHandle, h2Desc, h2_dev, h3Desc, h3_dev);
	bias2(cudnnHandle, h3Desc, h3_dev);
	bn2(cudnnHandle, h3Desc, h3_dev, h3_bn_dev);
	relu(cudnnHandle, h3Desc, h3_bn_dev);
	max_pooling_2d(cudnnHandle, h3Desc, h3_bn_dev, h4Desc, h4_dev);

	// fcl
	l3(cublasHandle, batch_size, h4_dev, h5_dev);
	bias3(cudnnHandle, h5Desc, h5_dev);
	relu(cudnnHandle, h5Desc, h5_dev);
	l4(cublasHandle, batch_size, h5_dev, y_dev);
	bias4(cudnnHandle, yDesc, y_dev);

	// output
	checkCudaErrors(cudaMemcpy(y, y_dev, sizeof(y_t), cudaMemcpyDeviceToHost));
}

cuDNNでBatchNormalizationの順伝播は、cudnnBatchNormalizationForwardInferenceで行う。
実際の処理は、layers.hのBatchNormalizationクラスのoperator()に実装している。
Chainerで学習したモデルから読み込んだ、gamma.npy、beta.npy、avg_mean.npy、avg_var.npyをcudnnBatchNormalizationForwardInferenceのパラメータに渡してやればよい。

cudnnBatchNormalizationForwardInferenceの2番目のパラメータのmodeには、CUDNN_BATCHNORM_SPATIALを指定する。
畳み込み層の直後でBatchNormalizationを使用する場合は、CUDNN_BATCHNORM_SPATIALを使用する。
ChainerのソースでもCUDNN_BATCHNORM_SPATIALが使用されていることを確認した。

cudnnBatchNormalizationForwardInferenceの最後の引数のepsの値は、Chainerのデフォルト値(2e-5)を固定で入力している。

実行結果

Chainerで推論した結果と完全に一致することを確認した。


ソースをGitHubで公開しました。
github.com