この記事はDeep Learningフレームワークざっくり紹介 Advent Calendarの第13日目の記事です。

deeplearn.jsとは

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

特徴としては

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

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

Teachable Machine

deeplearn.jsの面白さを体感するにはこのTeachable Machineという最近できたデモで遊んでみるのがよいと思います。

左の画像をデータに右の教師データへの対応付を学習させるというものです。パソコンにWeb Cameraがついていればすぐに試すことができます。これはWeb Browser上で動くdeeplearn.jsで学習を行っています。つまりクライアントサイドで訓練と予測両方ができるフレームワークとなっています。

Getting Started

自分でdeeplearn.jsを使ったアプリケーションを作る場合にはnpmあるいはyarnのパッケージとなっているので

$ npm install deeplearn
Or
$ yarn add deeplearn

でプロジェクトに組み込むことができます。またdeeplearn.jsにはすでに多くのdemoアプリが作られているのでこちらを手元で走らせてみるのがよいでしょう。

まずプロジェクトをcloneしてきます。

$ git clone https://github.com/PAIR-code/deeplearnjs.git
$ cd deeplearnjs

プロジェクトのビルドを行います。

$ yarn prep

デモアプリケーションはdemosディレクトリの下にあり、watch-demoというスクリプトで起動することができます。例えばmodel-builderというアプリケーションを起動するなら

$ ./scripts/watch-demo demos/model-builder

watch-demoはファイルの更新を監視して定期的にリコンパイルをするのでアプリの様子を見ながらコードの変更ができます。

ちなみにmodel-builderはMNISTやCIFAR 10といった画像識別のためのモデルをブラウザ上で作成して、それをダウンロードすることができます。demoの下にはオフィシャルサイトのExampleに含まれていないものも多いので試してみてください。

Architecture

どうやってWebGLでのHardware Accelerationを実現しているかを簡単に紹介したいと思います。

Deep dive into deeplearn.js from Kai Sasaki

deeplearn.jsは数値計算用のKernelを持ったmathというフレームワークを内部的に持っています。簡単な四則演算から、行列の掛け算、畳み込みなどの計算を行うコードが含まれています。このmathというフレームワークはインターフェースのみ提供していて実装はNDArrayMathCPUとNDArrayMathGPUに委ねられています。これらが裏側でCPUで動かすかWebGL経由でのGPUコードを動かすかを決めます。つまりアプリケーションのコードからはmathという共通化されたAPIを叩くだけなのでCPUで動かそうと、WebGLで動かそうとアプリケーションコードは変わらないことになります。例えば行列の掛け算を行うコードは下記のようになります。

// CPUの場合はこちら
const math = new NDArrayMathCPU();
// GPUの場合はこちら
// const math = new NDArrayMathGPU();

const matrixShape = [2, 3];  // 2 rows, 3 columns.
const matrix = Array2D.new(matrixShape, [10, 20, 30, 40, 50, 60]);
const vector = Array1D.new([0, 1, 2]);
const result = math.matrixTimesVector(matrix, vector);

console.log("result", await result.data());

与えるmathによってどこで計算を行うかを切り替えることができます。TensorFlow likeにSessionを使った計算もできます。

const graph = new Graph();

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


const a: Tensor = graph.variable('a', Scalar.new(Math.random()));
const b: Tensor = graph.variable('b', Scalar.new(Math.random()));
const c: Tensor = graph.variable('c', Scalar.new(Math.random()));

const order2: Tensor = graph.multiply(a, graph.square(x));
const order1: Tensor = graph.multiply(b, x);
const y: Tensor = graph.add(graph.add(order2, order1), c);

const yLabel: Tensor = graph.placeholder('y label', []);
const cost: Tensor = graph.meanSquaredCost(y, yLabel);

// Sessionに対して計算プラットフォーム用のmathを与えてあげる。
const math = new NDArrayMathGPU();
const session = new Session(graph, math);

NDArrayMathCPUの場合には愚直にTypeScriptで実装されていますが、NDArrayMathGPUにはGPUのメモリにバインドされたTextureデータを操作するFragment ShaderとしてKernelが実装されています。この記事では詳細は割愛しますが、WebGLを通したGPGPUとしてのコードを見ることができて面白いです。

Future Works

deeplearn.jsはまだリリースされたばかりで活発に開発が続けられています。最後に今後予定されている機能、改善を取り上げたいと思います。詳細はRoadmapに記載されています。

Automatic TensorFlow to deeplearn.js

今でもTensorFlowのモデル(GraphDef)をdeeplearn.jsにインポートすることは可能ですが、計算グラフ自体を自分でdeeplearn.jsのコードとして書く必要がありかなり煩雑です。GraphDefファイルを読んで自動で計算グラフを構築する機能が予定されています。

Decoupling NDArray from storage mechanism

deeplearn.jsの基本的なデータ型であるTensorは内部のデータ保持に更にNDArrayという型を用いています。NDArrayは本来は多次元配列を表すだけの抽象的なデータ構造として利用されるべきなのですが、deeplearn.jsではCPUとGPU上でのデータ構造と密結合してしまっています。これを疎結合にしてデータ構造を柔軟に扱えるようにします。特にWebGLTextureの情報ははずせるようにしたいです。

Model zoo

deeplearn.jsで使えるようにした学習済モデルを増やしていきます。すでにここに幾つかありますが、これを別管理にできるようにするのだと思います。

まとめ

この他にもかっこいい感じのdemoアプリはいつでもcontribution welcomeらしいので是非興味がありましたら、作ってみてください！

Future works of deeplearn.js