この記事はTensorFlow Advent Calendar13日目の記事です。

deeplearn.jsとは

deeplearn.jsはWeb Browser上で走るDeep Learningフレームワークです。今年の8月にGoogleのPAIRというプロジェクトから発表されました。

特徴としては

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

deeplearn.jsの紹介記事はすでに幾つか見つけたのですが、自分が開発に携わっているということもあり紹介したいと思います。

詳細はDeep Learningフレームワークざっくり紹介Advent Calendarにも記載しました。

Port TensorFlow model

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

なお詳細はここに記載されています。

Checkpointファイルの生成

Saving and Restoringに記載されているようにtf.train.Saverを使ってCheckpointファイルを生成します。deeplearn.jsのリポジトリにはdemos/mnist/fully_connected_feed.pyというお試しスクリプトがあるので、こちらで生成することができます。

Dump variables

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

また試してないのですが、スクリプト内のサポートしているモデルタイプにPyTorchがあるので、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として読めるようになっています。Float32ArrayはJavaScriptのTypedArrayのひとつで、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を使って先程のバイナリファイルをArrayBufferとして取得してきます。

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として結果が得られます。

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に"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のデモサイト

Port TensorFlow to deeplearn.js

deeplearn.jsとは

Port TensorFlow model

Checkpointファイルの生成

Dump variables

CheckpointLoader

問題

まとめ

Reference