Port TensorFlow to deeplearn.js

この記事は[TensorFlow Advent Calendar](https://qiita.com/advent-calendar/2017/tensorflow)13日目の記事です。


# deeplearn.jsとは

[deeplearn.js](https://deeplearnjs.org/)はWeb Browser上で走るDeep Learningフレームワークです。今年の8月にGoogleのPAIRというプロジェクトから[発表](https://research.googleblog.com/2017/08/harness-power-of-machine-learning-in.html)されました。

<img width="1273" alt="Screen Shot 2017-11-30 at 19.11.00.png" src="https://qiita-image-store.s3.amazonaws.com/0/7129/f9e41a42-96b6-7827-0c04-a23e9fdf0211.png">

特徴としては

* Web Browser上のJavaScriptコードとして動くブラウザアプリケーション用のフレームワーク
* WebGLを使ったHardware Accelerationが可能
* GoogleのPAIRというプロジェクトから[OSSとして公開されている](https://github.com/PAIR-code/deeplearnjs)

[deeplearn.jsの紹介記事](https://qiita.com/search?q=deeplearn.js)はすでに幾つか見つけたのですが、自分が開発に携わっているということもあり紹介したいと思います。

詳細は[Deep Learningフレームワークざっくり紹介Advent Calendar](https://qiita.com/advent-calendar/2017/deep-learning-frameworks)にも記載しました。


# Port TensorFlow model

deeplearn.jsにはTensorFlowのCheckpointファイルをインポートする機能があります。deeplearn.js単体でも訓練から推論までできるのですが、ブラウザで大規模な学習を行わせることが賢い選択と言えないことも多くあります。こういった場合に既存のモデルを利用できるととても便利で効率的です。TensorFlowのCheckpointファイル利用するためには幾つかの手順が必要となりますので、今回のdeeplearn.jsでTensorFlowの訓練済モデルを使う方法を記述いたします。

なお詳細は[ここ](https://deeplearnjs.org/demos/mnist/mnist.html)に記載されています。

## Checkpointファイルの生成

[Saving and Restoring](https://www.tensorflow.org/programmers_guide/saved_model)に記載されているように`tf.train.Saver`を使ってCheckpointファイルを生成します。deeplearn.jsのリポジトリには[`demos/mnist/fully_connected_feed.py
`](https://github.com/PAIR-code/deeplearnjs/blob/master/demos/mnist/fully_connected_feed.py)というお試しスクリプトがあるので、こちらで生成することができます。

## Dump variables

実はCheckpointファイルはそのままでは使えないので、必要な変数をdumpしてあげる必要があります。この処理は[`scripts/dump_checkpoints/dump_checkpoint_vars.py`](https://github.com/PAIR-code/deeplearnjs/tree/master/scripts/dump_checkpoints)を使って行うことができます。内部的には`tf.train.NewCheckpointReader`を使って読ませています。どうもCheckpointファイルのフォーマットというのは[undocumenteなもの](https://github.com/tensorflow/tensorflow/issues/3967)らしく、これがpublicなフォーマットとなれば、ダイレクトにCheckpointファイルを読むことも可能になるかもしれません。

また試してないのですが、スクリプト内のサポートしているモデルタイプに[PyTorch](https://github.com/PAIR-code/deeplearnjs/blob/master/scripts/dump_checkpoints/pytorch_checkpoint_dumper.py)があるので、PyTorchのモデルを利用できるのだと思います。

このスクリプトを走らせると各変数毎のバイナリファイルができます。

```
$ tree demos/mnist
demos/mnist
├── fully_connected_feed.py
├── hidden1_biases
├── hidden1_weights
├── hidden2_biases
├── hidden2_weights
├── index.html
├── manifest.json
├── mnist.md
├── mnist.ts
├── sample_data.json
├── softmax_linear_biases
└── softmax_linear_weights
```

このファイルはJavaScriptの[`Float32Array`](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Float32Array)として読めるようになっています。Float32ArrayはJavaScriptのTypedArrayのひとつで、[ArrayBuffer](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/ArrayBuffer)に対して32ビット浮動小数点としてのアクセス方法を提供するViewです。ArrayBufferを使えば高速にJavaScriptから物理メモリにアクセスできるため、deeplearn.jsの基本的なデータ型であるNDArrayのバックエンドにも使われています。

## CheckpointLoader

この変数を読むためにはdeeplearn.jsが提供する`CheckpointLoader`を使います。`CheckpointLoader`はまず`manifest.json`という、変数を格納したファイル名とそのTensorの型を記述したファイルを読みます。先程の例だと下記のような`manifest.json`ができているはずです。

```
{
  "hidden1/biases": {
    "filename": "hidden1_biases",
    "shape": [
      128
    ]
  },
  "hidden1/weights": {
    "filename": "hidden1_weights",
    "shape": [
      784,
      128
    ]
  },
  "hidden2/biases": {
    "filename": "hidden2_biases",
    "shape": [
      32
    ]
  },
  "hidden2/weights": {
    "filename": "hidden2_weights",
    "shape": [
      128,
      32
    ]
  },
  "softmax_linear/biases": {
    "filename": "softmax_linear_biases",
    "shape": [
      10
    ]
  },
  "softmax_linear/weights": {
    "filename": "softmax_linear_weights",
    "shape": [
      32,
      10
    ]
  }
}
```

`CheckpointLoader`はこのファイルから読むべきTensorとその大きさをしることができます。次に[`XMLHttpRequest`](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest)を使って先程のバイナリファイルを`ArrayBuffer`として取得してきます。

```typescript
const xhr = new XMLHttpRequest();
// ArrayBufferとして結果を取得
// see: https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType
xhr.responseType = 'arraybuffer';
const fname = this.checkpointManifest[varName].filename;
xhr.open('GET', this.urlPath + fname);
xhr.onload = () => {
  const values = new Float32Array(xhr.response);
  const ndarray =
    NDArray.make(this.checkpointManifest[varName].shape, {values});
  resolve(ndarray);
};
```

つまり`XMLHttpRequest`から読めるようなところに予めモデルファイルを配置しておく必要があります。ロードが終わるとTensor名とデータのmapとして結果が得られます。

```typescript
const varLoader = new CheckpointLoader('.');
varLoader.getAllVariables().then(async vars => {
  const math = new NDArrayMathGPU();

  // manifest.jsonで設定されたTensor名とデータのmapが得られる。
  const hidden1W = vars['hidden1/weights'] as Array2D;
  const hidden1B = vars['hidden1/biases'] as Array1D;  
  // ...

  // 得られたものはNDArrayなのでdeeplearn.jsのoperatorに使うことができる
  const hidden1 =
      math.relu(math.add(math.vectorTimesMatrix(..., hidden1W), hidden1B)) as
      Array1D;
  const hidden2 =
      math.relu(math.add(
          math.vectorTimesMatrix(hidden1, hidden2W), hidden2B)) as Array1D;

  const logits = math.add(math.vectorTimesMatrix(hidden2, softmaxW), softmaxB);

  const label = math.argMax(logits);

  console.log('Predicted label: ', await label.data());
});
```

これでdeeplearn.js側でモデルを構築することができます。

# 問題

ただこの方法だと幾つか面倒な点があります。

* Checkpointファイルをdeeplearn.jsアプリケーションから直接読めないためArrayBufferとして読めるようにファイルに分割しないといけない
* CheckpointLoaderは単なるTensorのコレクションなので、TensorFlowで使っていたOpsの情報を持っていません。そのため再利用するためにはdeeplearn.jsの方で完全な計算グラフを再構築する必要がある。`GraphDef`などからその情報を得られればできるかもしれないが、まだTensorFlowにはあってdeeplearn.jsにないOpsが多くあるので現状難しい

というわけでこういった問題を解決するために[Roadmap](https://deeplearnjs.org/docs/roadmap.html)に"Automatic TensorFlow to deeplearn.js"というタスクがつまれています。


> Currently we support dumping weights from a TensorFlow checkpoint into a format that can be imported into deeplearn.js, however the
> developer must then recreate the model in deeplearn.js and use the weights from that checkpoint.
> 
> We plan on building a way to port models directly from TensorFlow to deeplearn.js automatically from a GraphDef.

# まとめ

deeplearn.jsはGoogleのエンジニアが20%ルールで作っているプロダクトらしいですが、活発に開発が続けられています。とりわけ、deeplearn.jsを使ったデモは作るのも使うのも面白いので是非試してみてください！

[deeplearn.jsのデモサイト](https://deeplearnjs.org/index.html#demos)

# Reference

* [deeplearn.js](https://deeplearnjs.org/)
* [tf.train.NewCheckpointReader](https://www.tensorflow.org/api_docs/python/tf/train/NewCheckpointReader)
* [XMLHttpRequest](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest)
* [ArrayBuffer](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/ArrayBuffer)
* [Float32Array](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Float32Array)