JCuda - ニューラルネットの実装

ポイントは２つあります。

1. 出力ノード別に並列化する
2. BATCHSIZEに展開して並列化する

さて、ニューラルネットのプログラムは以下のソースファイルから構成されるのでした。

• SimpleNet.java
• NeruralNet.java
• ILoader.java
• ImageLabelSet.java
• MNIST.java
• ILossFunction.java
• QuadLossFunction
• TestMnist.java

JCuda の初期化メソッド

public class NNUtil {
  public static Map<String, CUfunction> initJCuda(String cuFileName) {
        // 例外処理を有効にする
        JCudaDriver.setExceptionsEnabled(true);

                // cu ファイルから ptx ファイルを（必要なら）コンパイルしてロー>ド
        String ptxFileName = preparePtxFile(cuFileName);

        // CUDA ドライバーを初期化し、最初のデバイスに対するコンテキストを作成する
        cuInit(0);
        CUcontext pctx = new CUcontext();
        CUdevice dev = new CUdevice();
        cuDeviceGet(dev, 0);
        cuCtxCreate(pctx, 0, dev);

        // PTX ファイルを読み込む
        CUmodule module = new CUmodule();
        cuModuleLoad(module, ptxFileName);

        // カーネル関数を展開する
        Map<String, CUfunction> fMapper = new ConcurrentHashMap<String, CUfunction>();
        createMapper(fMapper, module);
        return fMapper;
  }
}

MNIST.java へのデバイスメモリの追加

このとき、2次元配列の要素 x, y へのアクセスは devMem[x][y] で表されます。
devMem[x] で対象となる devMemArray を選択し、この devMemArray を使って devMemArar[y] を選択するイメージとなります。

インスタンスを作成した後に画像データを指定できるようにするため、setContentDev() メソッドを用意しました。クラスによってシグニチャが多少異なりますが、一連の操作の前に必ずこのメソッドを呼ぶ必要があります。

今回は MNISTImage, MNISTLabel 内でデバイスメモリを管理するようにしました。
ILoad にメモリ管理用の3つのメソッドを追加します。

• CUdeviceptr getContentDev(int index)

  指定した index に対応する devMemArray へのポインタを返します。

• CUdeviceptr getContentDev()

  devMem へのポインタを返します。

• void setContentDev()

  デバイスメモリにホストメモリの内容を転送する。

SimpleNet へのデバイスメモリの追加

パーセプトロンへのデバイスメモリの追加

そのために devOutz, devOutzArray, devTmpz, devTmpzArray を定義します。
これらはパーセプトロンの出力値、線形和の結果保持に使用されます。

batchsize に関して並列化するため、入力、出力ともに2次元配列として定義されることに注意します。
線形和や非線形関数は batchsize 方向の添字によらない値となるため、3重ポインタは必要ないことにも注意してください。

最終段は softmax 関数となるため、出力値が全て揃わないと計算できません。
そのため、__syncthreads() で全てのスレッドの計算終了を待ってから非線形関数の計算を行っています。
shared memory を使うことも検討しましたが、64k bytes の制限に収まらなかったので、今回は採用を見送りました。
shared memory を使うともっと速くなる可能性があります。

forward メソッドを2次元のデバイスメモリを入力とし、2次元のデバイスメモリを出力とするメソッドに書き換えます。出力段は常に devOutz が返されます。

CUdeviceptr forward(CUdeviceptr in) {
  int fmt = format2no(format);
  Pointer kp = Pointer.to(
          Pointer.to(devOutz), Pointer.to(devTmpz),
          Pointer.to(devW), Pointer.to(devB), Pointer.to(in),
          Pointer.to(new int[]{inn}), Pointer.to(new int[]{outn}),
          Pointer.to(new int[]{fmt}), Pointer.to(new int[]{batchsize})
          );
  cuLaunchKernel(fMapper.get("calc_forward"),
          calcBlock2D(outn), calcBlock2D(batchsize), 1,
          NTHREAD2, NTHREAD2, 1,
          0, null,
          kp, null
          );
  cuCtxSynchronize();
  return devOutz;
}

extern "C" __global__ void
calc_forward(float **out, float **tmp, float **w, float *b, float **in, int inn, int outn, int fmt, int batchsize) {
  const int o = blockDim.x * blockIdx.x + threadIdx.x;
  const int s = blockDim.y * blockIdx.y + threadIdx.y;

  // 線形和の計算
  if (o < outn && s < batchsize) {
    float t = 0.0f;
    for (int i = 0; i < inn; ++i) {
      t += in[s][i] * w[i][o];
    }
    t += b[o];
    tmp[s][o] = t;
  }
  __syncthreads();

  // 非線型関数の出力
  if (o < outn && s < batchsize) {
    if (fmt == 1) {
      // softmax 関数
      float xmax = calc_xmax(tmp[s], outn);
      float divisor = calc_div(tmp[s], outn, xmax);
      outz[s][o] = calc_softmax(tmp[s], xmax, divisor, o);
    } else {
      // sigmoid 関数
      outz[s][o] = sigmoid(tmp[z][o]);
    }
  }
}

逆誤差散乱法記述の修正

calc_deriv_b は bMostouter が true の時と false の時で引数が異なります。

• true の時

  引数並びは calc_b_deriv(devDB, devBhot, devSamples, true) となります。
  devSamples は連番を乱数に変換するための、batchsize 個の要素を持つ int 配列です。
  samples は host 側で作成し、予めデバイスメモリ devSamples にコピーしておく必要があります。

• false の時

  引数並びは calc_b_deriv(devDB, devW, devOutz, devDB2, false) となります。
  devDB2 は1つ下の層の devDB を入力します。

        public void calc_deriv_w(CUdeviceptr devDW, 
                CUdeviceptr in, int xsize, CUdeviceptr delta, int ysize) {
                Pointer kp = Pointer.to(Pointer.to(devDW),
                                Pointer.to(in), Pointer.to(new int[]{xsize}),
                                Pointer.to(delta), Pointer.to(new int[]{ysize}),
                                Pointer.to(new int[]{batchsize})
                );
                cuLaunchKernel(fMapper.get("calc_deriv_w_kernel"),
                                calcBlock2D(xsize), calcBlock2D(ysize), 1,
                                NTHREAD2, NTHREAD2, 1,
                                0, null,
                                kp, null);
                cuCtxSynchronize();
        }

// 損失関数の w 方向の微分
extern "C"
__global__ void calc_deriv_w_kernel(float **dw, float **in, int xsize, float **d
b,
  int ysize, int bs) {
  const int x = blockDim.x * blockIdx.x + threadIdx.x;
  const int y = blockDim.y * blockIdx.y + threadIdx.y;
  
  if (x < xsize && y < ysize) {
    float d = 0.0f;

    // batchsize 方向に加算してしまう
    for (int z = 0; z < bs; ++z) {
      d += in[z][x] * db[z][y];
    }
    dw[x][y] = d;
  }
}

確率的最急降下法の適用

// b に関する学習
extern "C" __global__ void
learn_1d(float *bout, floar **deriv. float lrate, int size, int nsample, int bs) {
  const int i = blockDim.x * blockIdx.x + threadIdx.x;
  
  if (i < size) {
    for (int s = 0; s < bs; ++s) {
      bout[i] -= lrate * deriv[s][i] / nsample; 
    }
  }
}

デバイスメモリの解放

具体的にはデバイスメモリを確保したクラスに public void finalize() を作成し、後処理を記述していきます。
SimpleNet の場合の記述は以下のようになります。

        public void finalize() {
                IntStream.range(0, inn).forEach(i->{
                        cuMemFree(devWArray[i]);
                });
                cuMemFree(devW);        
                cuMemFree(devB);
                IntStream.range(0, batchsize).forEach(i->{
                        cuMemFree(devOutz2D[i]);
                });
                cuMemFree(devOutz);
                IntStream.range(0, batchsize).forEach(i->{
                        cuMemFree(devTmpz2D[i]);
                });
                cuMemFree(devTmpz);
        }