これはTensorFlow 2.0 Advent Calendar 12日目の記事です。

機械学習用のフレームワークとしてTensorFlowは欠かせないツールのひとつとなっています。今年は2.0がリリースされ、TensorFlowは大きな進化を経た年となりました。その中でも最近注目を集めているのがTensorFlow.jsです。プロジェクトしては小さく始まったTensorFlow.jsですが、今年は1.0がリリースされ、多くの企業や個人が機械学習のアプリケーションをTensorFlow.js上で書くようになっています。

この記事では今年ますます活発に開発が行われたTensorFlow.jsの新機能のいくつかを紹介します。

TensorFlow.jsとは

TensorFlow.jsはTensorFlowのJavaScript拡張で、JavaScriptからTensorFlowの機能を利用することができます。TensorFlowは他にもJavaやGo、C/C++から利用することができますが、JavaScriptはPythonに次いで活発に開発がされているAPIとなっています。

- https://www.tensorflow.org/js/

元々はdeeplearn.jsとしてGoogleの20%ルールのプロジェクトとして始められたようですが、去年、TensorFlowのプロジェクトに統合されました。

私もcontributorとして開発に携わってきて、非常に活気ある興味深いプロジェクトのひとつだと感じています。

TensorFlow.jsの特徴

TensorFlow.jsの大きなゴールはWebプラットフォームを機械学習基盤としてもっと利用しやすくすることです。WebはWeb Browserというインタラクティブ性に優れたプラットフォームをすでにもっており、この数十年で多くのアプリケーションを生み出してきました。機械学習を用いたアプリケーションを作る上でこのプラットフォームを利用することはWebの可能性を広げるだけでなく、機械学習の利用価値を広く知らしめることにも繋がります。

またWebのデファクトスタンダードとなっているJavaScriptは最も利用されているプログラミング言語であり、JavaScriptで利用できるということはこれらのユーザ、開発者と共に開発を進めていくことができるということを意味します。

こういった機械学習の裾野を広げるという目的を掲げTensorFlow.jsは開発されています。

また最近ではプライバシー保護の観点から個人に紐づく情報や行動ログをサーバに送信したくないといった人もでてきました。こういった場合でも機械学習のソリューションを提供するために、クライアントサイドで機械学習モデルを訓練して完結させるといった考え方がでてきました。Edge AIやFederated Learningといった考え方も同様にクライアントサイドでの機械学習を後押しする考え方として最近叫ばれており、TensorFlow.jsはこの方向性を実現するソリューションのひとつとなるでしょう。

それでは今年導入されたパフォーマンスや他のプラットフォームへの親和性を強化する機能たちを紹介していきます。

New Backends

TensorFlow.jsでは計算グラフのノードを実行するのにkernelという単位で実行しています。例えば

- 2つのtensorを足し合わせる add
- tensorのメタデータを変えるreshape
- 畳み込みを行う conv

これらの計算ノードをグラフ上に構成することでモデルとなる計算グラフを構成します。この中でモデルパラメタを更新することでモデルの訓練を実現しています。

さてこのモデルの実行を早くするには何が必要でしょうか。大きく分けて２つのやり方があります。

1. 計算グラフの構造を最適化する
2. ひとつひとつの計算ノードの実行を早くする

1. の方法では不必要な計算ノードを取り除いたり（pruning）、精度の低い浮動小数点（quantization）を利用することでより計算速度の速いモデルを作成します。これらはモデルのFreezingなどの手法をつかって実現することができます。

2. では計算ノードの実行コードを高速にすることでモデルの実行を早くします。GPUなど、実行基盤の異なる環境でモデルを走らせる場合にはこちらの手法をとっています。そのため、ひとつひとつの計算ノードのoperationに対してCPU、GPUなど実行環境に対応する実装を作る必要があります。逆にいえばハイレベルな計算グラフのレイヤーでは透過的に計算ノードのoperationを利用できるため、一度実装が作られてしまえば、ユーザは簡単にGPUなどの計算資源を利用することができます。

TensorFlow.jsでも同じようにbackendと呼ばれるレイヤーが実行環境の詳細を隠蔽するため一度TensorFlow.jsのモデルを構成すれば、様々な環境でモデルを実行することができます。

現在TensorFlow.jsでは下記のbackendがサポートされています。

- CPU
- WebGL
- WebGPU
- WebAssembly
- Node.js

ここではWebAssemblyに触れます。

WebAssembly Backend

WebAssemblyはWebプラットフォーム上のStackベースの仮想マシンで動く命令セットの標準です。WebAssemblyは仮想マシンと実ハードウェアの橋渡しをして、よりネイティブコードに近いパフォーマンスを実現するようにデザインされています。tfjs-backend-wasmというパッケージはWebAssembly用のbackend実装をTensorFlow.jsに提供しています。

利用は下記のようにimportしたtfjs-backend-wasmパッケージからtf.registerBackendでruntimeに登録をします。tf.setBackendで"wasm"に切り替えれば準備完了です。あとはTensorFlow.jsのCore APIやLayers APIでモデルを走らせるだけです。

import * as tf from '@tensorflow/tfjs';
import * as BackendWasm from 'tfjs-backend-wasm';
import {BackendWasmModule} from 'tfjs-backend-wasm';


tf.registerBackend('wasm', async () => {
  return new BackendWasm(new BackendWasmModule());
}, 3 /*priority*/);
await tf.setBackend("wasm");
const backendName = tf.getBackend();
console.log(backendName);

const t1 = tf.tensor([1,2,3]);
t1.print();

WebAssembly backendはalpha版という扱いなのでAPIや使い方は変わる可能性がありますが、よりネイティブに近いパフォーマンスがTensorFlow.jsで得られるとしたら魅力的です。

それではパフォーマンス向上ではなく、TensorFlow.jsの実行環境を拡張するパッケージを次に紹介します。

React Native

React NativeはiOSやAndroidなどのネイティブアプリを作成するためのプラットフォームで、JSXというJavaScriptの拡張を使ってアプリを書きます。 そのため、JavaScriptにコンパイルされるTensorFlow.jsのコードを実行することもできるのです。tfjs-react-nativeはこのReact Native上でTensorFlow.jsのモデルを動かすためのパッケージです。モバイル領域での機械学習ライブラリとしてはTensorFlow Liteが人気ですが、tfjs-react-nativeを利用することでも簡単にTensorFlowのモデルをネイティブアプリ上で動かすことができます。

tfjs-react-nativeパッケージをインポートしたあと下記のようなコード書くだけでTensorFlow.jsがReact Native上で動きます。

import React from 'react';
import {Component} from 'react';
import { StyleSheet, Text, View } from 'react-native';
import * as tf from '@tensorflow/tfjs';
import * as tfjsNative from '@tensorflow/tfjs-react-native';

export default class TensorFlowJSApp extends Component {
 private aStr: string;
 private bStr: string;
 private cStr: string;
 constructor(props) {
   super(props);
   const a = tf.tensor([[1, 2], [3, 4]]);
   const b = tf.tensor([[1, 2], [3, 4]]);

   const c = a.add(b);

   this.aStr = a.toString();
   this.bStr = b.toString();
   this.cStr = c.toString();
 }

 render() {
   return (
     <View style={styles.container}>
       <Text>Hello, TensorFlow.js! {'\n'}{'\n'}{this.aStr} + {'\n'}{this.bStr} = {'\n'}{this.cStr}</Text>
     </View>
   );
 }
}

下記はiOSシミュレータ上で動かした様子です。計算結果が正しく表示されています。

React Native自身もモバイルデバイスのプラットフォームとなっているので、React Nativeで動いたモデルはReact Nativeでサポートされているプラットフォーム（e.g. iOS, Android) でもそのまま動くようになるのが嬉しいですね。

Electron

同じようなネイティブアプリを作るプラットフォームとしてはElectronがあります。こちらはPC用のネイティブアプリを作るためのフレームワークで元々はGitHubが開発してオープンソースソフトウェアとして公開しました。SlackのデスクトップアプリなどもこのElectronで作られています。

ElectronのアプリはNode.js上で動くJavaScriptアプリケーションとして作られています。TensorFlow.jsはNode.js経由でTensorFlowのC APIを叩くbackendを持っているため、これを使ってTensorFlow.jsのアプリケーションをElectron上と共に動かすことができます。下記のように試してみました。

- tfjs-electron-app

こちらのコマンドでアプリが立ち上がります。

$ yarn && yarn start

正しく計算結果が表示されています。

ElectronもOSに依存しない抽象化されたレイヤーをアプリケーションに提供されているため、TensorFlow.jsで書いたモデルはmacOS, WindowsそしてLinux上でネイティブアプリに組み込んで利用することができます。

AutoML Integration

最後に紹介するのがGCPのCloud AutoMLのモデルをダイレクトに使えるようにしたパッケージ、tfjs-automlです。Cloud AutoMLは必要なデータを与えるとGCPの計算資源を使い、モデルのチューニング、訓練を行ってくれるクラウドサービスです。

TensorFlow.jsはモデルの訓練を行うこともできるのですが、大規模なデータを使う場合にはこのようなクラウドサービスを利用するのがお手軽です。最初にサーバサイドで訓練をして汎用的なモデルを作ったら、クライアントサイドでローカルなデータに基づきFine Tuningを行うというのが効率的なやり方かもしれません。

下記のようにモデルをロードすることができます。

import * as automl from '@tensorflow/tfjs-automl';
const modelUrl = 'model.json'; // URL to the model.json file.
const model = await automl.loadImageClassification(modelUrl);

TensorFlow.jsではtfjs-converterというツールを使ってKerasのHDF5モデルやSavedModel、freezedされたモデルをTensorFlow.jsで利用可能なフォーマットに変換できるのですが、tfjs-automlを使えば、Auto MLで生成したモデルに関してはダイレクトに持ってくることができます。

最後に

TensorFlow.jsはまだまだ新しいライブラリですが活発に開発が行われていて、これからも面白い機能が続々でてくるはずです。

ここで紹介した内容は先日出版された"Hands-On Machine Learning with TensorFlow.js" にも記載されています。TensorFlow.jsを使って様々な機械学習モデルを作成するHands-Onとなっています。ご興味ありましたらぜひ読んでみてください。

コード

画像