はじめてのdeeplearn.js

Googleが発表したdeeplearn.jsで、学習から推論までを試してみます。

個人的に一番グッときたのはWebGLを介したGPUの利用です。過去のディープラーニングライブラリはたいていCUDAをバックエンドに利用しており、nVidiaの比較的新しいGPUを積んだマシンでしかGPUの恩恵が得られません。私が主に利用しているマシンはMacBook Pro(13inch Mid2012)ですが、Intel HD Graphics 4000 というオンチップGPUが載っています。nVidiaには全然劣りますが数百GFLOPSは出るようですので、学習も多少は速くなるのでは、という期待があります。

公式のチュートリアル等から見る感じですと、TensorFlowで学習をしたモデルをブラウザで推論に使う、というような使い方が普通かと思いますが（当然、JavaScriptにしたモチベはそこなのでしょうし。。）、上記の理由から、主に学習周りを調べてみます。

環境設定

ビルドはインストラクション通りです。（下記のような使い方ではbuild-standaloneは不要かもです。）

git clone https://github.com/PAIR-code/deeplearnjs.git
cd deeplearnjs
npm run prep
./scripts/build-standalone.sh

最短で使うには公式のデモまわりの仕組みを利用するのが早そうです。demos/mydemoとか適当なフォルダを作ってTypeScriptのプログラム本体test.tsと表示用のindex.htmlを置きます。ちゃらっと試したい方はこちらをコピーしてください。

表示用のhtmlはこんなので十分です。

<h1>deeplearn.js trial</h1>
<div><span id="output"></span></div>
<script src="bundle.js"></script>

test.tsはこちらです。中身は次章で説明します。これがbundle.jsにビルドされます。

この状態で、deeplearn.jsのベースディレクトリから、

 ./scripts/watch-demo demos/mydemo/mydemo.ts

でサーバが起動します。tsファイルを定期監視して変更があればビルドを行ってくれます。ブラウザから

http://localhost:8080/demos/mydemo/

でアクセスできます。console.log()の出力結果を見るためにブラウザのJavaScriptコンソールを開きます。

学習から推論まで

test.tsの中身は、関数近似で学ぶ chainer とディープラーニングのお題をdeeplearn.jsでやってみたものです。$y=e^x$ をMLPで近似します。いまのところここぐらいしか情報がないのでほとんどコピペです。

バッチの入力場所はgraph.placeholderで、最適化される係数はgraph.variableで生成します。TensorFlow的ですね。初期値はNDArray(2DならArray2D)で渡します。numpyを意識したのでしょうね。

const x: Tensor = graph.placeholder('x', []);

const W1data = Array2D.randNormal([16, 1]);
const b1data = Array2D.zeros([16, 1]);

const W1: Tensor = graph.variable('W1', W1data);
const b1: Tensor = graph.variable('b1', b1data);

グラフを構成していきます。最後のreshapeは、ラベルと次元をあわせるために1x1行列からconstantに変換するために入れています。ワークフロー的にもTensorFlowにかなり近いですね。

const h1: Tensor = leaky_relu(graph.add(graph.multiply(W1, x), b1));
const h2: Tensor = leaky_relu(graph.add(graph.matmul(W2, h1), b2));
const y_: Tensor = leaky_relu(graph.add(graph.matmul(W3, h2), b3));
const y: Tensor = graph.reshape(y_,[]);

leaky_reluが無かったので手書きします。イマイチですが。。

function leaky_relu(x: Tensor): Tensor {
  const mx: Tensor = graph.multiply(graph.constant(-1), graph.relu(x));
  const neg_part: Tensor = graph.multiply(graph.constant(-0.2), graph.relu(mx));
  return graph.add(graph.relu(x) , neg_part)
}

バッチデータの供給はInCPUMemoryShuffledInputProviderというひとを使います。その名の通り、session.train()から呼ばれて、placeholderに対してデータを供給するひとのようです。

  for (var i = 0; i < 100; i++) {
    var xr = Math.random();
    xs.push(track(Scalar.new(xr)));
    ys.push(track(Scalar.new(Math.exp(xr))));
  }

  const shuffledInputProviderBuilder =
      new InCPUMemoryShuffledInputProviderBuilder([xs, ys]);
      const [xProvider, yProvider] =
      shuffledInputProviderBuilder.getInputProviders();

CPU利用かGPU利用かはNDArrayMathCPUを使うかNDArrayMathGPUを使うかで切り替えできます。下記はCPUの例。tsconfig.jsでnoUnusedLocals=trueと設定されているため、import文も書き変えが必要ですので注意です。

const math = new NDArrayMathCPU();
const session = new Session(graph, math);

math.scope((keep, track) => {
..
}

訓練はsession.train()です。SGDしか使えないそうです。sessionスコープ内でクリアされたくないデータはtrackで囲むそうです。正直よく理解できていません。そもそもJSでのスコープが理解できておらず、変数が軒並みconst定義してあるのですがそういうものなのでしょうか。

  for (let i = 0; i < NUM_BATCHES; i++) {
      const costValue = session.train(
        cost,
        [{ tensor: x, data: xProvider }, { tensor: yLabel, data: yProvider }],
        BATCH_SIZE, optimizer, CostReduction.MEAN);
        var cost_val = costValue.get();
        console.log('average cost: ' + cost_val);
    }

ちなみにサンプルは下記が間違ってるんじゃないですかね。yじゃなくてyLabelに供給しないと。

        [{ tensor: x, data: xProvider }, { tensor: y, data: yProvider }],       // NG
        [{ tensor: x, data: xProvider }, { tensor: yLabel, data: yProvider }], // OK

推論はsession.eval()します。グラフと、placeholderと対応するデータを与えます。返ってきたテンソルをgetValue()して値を取り出します。下記の例では、xに0.2を代入したときの値を出します。

var hogehoge = session.eval(y, [{tensor: x, data: track(Scalar.new(0.2))}]).getValues();

下記が実行結果のスクリーンショットです。

この例だともうちょっと学習は回したほうがよさそうですね。

なんだかたまに（というか結構な頻度で）収束に失敗します。係数が変化しなくなります。微分消失なのか、なんらかのバグなのか。
乱数のseedを固定する方法もないし、ちょっと簡単に確認できないので目をつぶります。

ベンチマーク

CPU/GPU差とchainerとの比較を調べます。chainer側をなるべく同条件に合わせます。どうもミニバッチサイズが１に固定ということらしい？のでchainer版もそういうループにしました。ちなみに今回の問題設定ではミニバッチを少なくしてひたすら回数を回すほうが収束性はよいです（データにノイズがありませんし）。Chainer版のコードはこちら。

ネットワークの各レイヤの出力ノード数が16ch,32ch の場合のベンチマーク結果です。

	Train [sec]	Prediction[msec/cycle]
deeplearning.js(GPU):	30.31	3.02
deeplearning.js(CPU):	5.36	0.23
Chainer(CPU):	14.49	0.55

おっと・・直感と反した値が出ました。。CPUのほうが速いし、なんとChainerより速い。たかだか16x32の行列演算ではオーバヘッドのほうがデカイということでしょう。

ということで行列演算の規模を大きくします。256x1024 にします。

	Train [sec]	Prediction[msec/cycle]
deeplearning.js(GPU):	31.63	4.12
deeplearning.js(CPU):	352.22	10.42
Chainer(CPU):	28.62	0.65

なるほど。行列サイズが大きくなっても、GPUはほとんど差がありません。一方でCPUはかなり遅くなりますね。ただしこの条件だと、なんだかいずれの場合もdeeplearn.jsでは収束しませんでしたので怪しいです。

Chainerとのベンチマークも、フェアな比較なのかちょっと不明ですね。学習周りとか何も読んでないのでなんともです。ただ、学習時間比では明らかな違いが出ていますので、演算量が増えれば増えるほどGPUが有利になるのは想像できます。

まあ、例もよろしくないですし、あくまで参考ということでお願いいたします。

まとめ

ひとまず学習から推論までの流れと、WebGL利用のGPU対応の可能性については確認できました。ベンチマークに線形結合を用いたのがちょっと失敗だったと思います。CNNとかだと差が出やすいでしょうか。学習周りの実装はまだまだこれからだと思いますが、楽しみです。

TypeScriptなりWebプログラミングが不慣れなので苦労しました。新しいライブラリはREPLなどで試行錯誤しながら習得したいところですが、ブラウザ側のWebGL実装を使うという仕組み上nodeのREPLなどで対応するのも難しそうです。一方で、TensorFlowと非常に近い作りになっている（Keras風な書き方もできるようです）のと、TypeScriptを書けばビルドやサーバを立てるのは全部スクリプトでやってくれるので、サンプルのマネをするレベルであれば、思ったより敷居は高くありませんでした。

デプロイ用のスクリプトも入っているようで、推論のデモなどに使うにはよい環境だと思います。今後の発展が楽しみです。