Kerasで構築したモデルをTensorFlow.jsで動かす

TensorFlow/Kerasで構築して学習したモデルをTensorFlow.jsで動かしてみた。

あらまし

書籍「AlphaZero 人工知能プログラミング実践入門」に沿って各ゲームの強化学習のモデルをGoogle Colab上で構築・学習させることができる。学習した結果をColab上で動かすことができるが、 Tensorflow.js（以降TFJS）を使ったらブラウザでも動かせるんじゃないかと思って見てみた。

モデルデータのコンバート（TensorFlow1.x系）

本家のTensorFlow/Kerasで構築・学習・保存したモデル、.h5ファイルをTFJSで動かすためにはコンバートする必要がある。「Keras モデルを TensorFlow.js にインポートする」では pip で tensorflowjs_converter をインストールしているが、TensorFlow1.x系に対応するバージョンを使用するにはどうすればいいか。

pipにも tensorflowjs があり、そちらをインストールすることでコンバータも使えるっぽい。そうした場合、バージョン２未満指定 pip install 'tensorflowjs<2.0 でよいっぽい。

コンバータもローカルじゃなくてColab上で動かした。 .h5ファイルをGoogle Driveに置いて実行：

# TensorFlow 1.x
%tensorflow_version 1.x
!pip uninstall -y tensorflow
!pip install tensorflow==1.15

import tensorflow as tf
print(tf.__version__)

# tensorflow_converter
!pip install 'tensorflowjs<2.0'

# Google Driveをマウント
from google.colab import drive
drive.mount('./drive')

# モデルファイルのコンバート（パスはよしなに）
!tensorflowjs_converter --input_format keras ./drive/MyDrive/AlphaZeroBook/Reversi/best.h5 ./drive/MyDrive/AlphaZeroBook/Reversi/tfjsmodel

コンバート結果がGoogle Driveに出力されるので、それをローカルにダウンロードする。

TensorFlow.js で動かす

コンバートしたモデルをTFJSで動かす。行う計算は、学習済みのモデルを使って推論するだけで、学習はしない。

TFJSを動かすには npm の @tensorflow/tfjs モジュールを使う。バージョンはこちらも <2.0 でよい。

以下、TFJSで必要ないくつかの操作をメモ：

モデルの読み込み

tf.loadLayersModel を使う：

const tf = require('@tensorflow/tfjs')

async function main() {
  // load model
  const path = 'http://localhost:8080/tfjsmodel/model.json'
  const model = await tf.loadLayersModel(path)
  ...

非同期なので await で待つ
node.js で実行する場合でもローカルファイルからは読み込めず、ローカルサーバ経由する必要がある

推論

モデルを使って推論するには、Pythonと同様に model.predict(...) を使う。

predict に与える引数や結果はテンソルなので、TFJS上でのテンソルの扱いを知る必要がある。

テンソルの操作

本家のPythonではnumpyのarrayを用いるが、TFJSではテンソルを扱う。

処理内容	記述
JS配列からテンソルへの変換	`let x = tf.tensor([[1, 2, 3], [4, 5, 6]])`
リシェイプ	`x = x.reshape([3, 2])`
推論	`const y = model.predict(x, 1)`
値の取り出し（JS配列化）	`const arr = y.arraySync()`

テンソルは [] で値を取り出せず、 arraySync() で配列に変換しないといけなくて結構面倒。

他に、Python・numpyで用意されているさまざまな関数は、TFJSでは大抵tf経由で使えるようになっている：

sum, div, argMax, zeros, …

ゲーム部分

Colabでの学習時にPythonで記述していたゲーム部分は、そっくりJavaScriptに持ってくる必要がある。この辺、二度手間なのでなんとかしたい感はある…。

node.js で動かす

ブラウザで動かすのもいいが、node.jsを使ってターミナル上で動かすのも手軽でよい。その場合は @tensorflow/tfjs-node か -gpu付き版をインポートに加えるとのこと。しかし自分の環境だと、 tfjs-node をインポートすると reshape でエラーが出てしまうため、動かせなかった：

TypeError: backend.reshape is not a function

インポートしなければ動くが、ワーニングが出る：

============================
Hi there 👋. Looks like you are running TensorFlow.js in Node.js. To speed
things up dramatically, install our node backend, which binds to TensorFlow C++,
by running npm i @tensorflow/tfjs-node, or npm i @tensorflow/tfjs-node-gpu if
you have CUDA. Then call require('@tensorflow/tfjs-node'); (-gpu suffix for
CUDA) at the start of your program.
Visit https://github.com/tensorflow/tfjs-node for more details.
============================

ワーニングメッセージを抑制するには環境 IS_NODE を false にしてやる：

tf.env().set('IS_NODE', false)

ソース

参考

Keras モデルを TensorFlow.js にインポートする
プラットフォームと環境 | TensorFlow.js
TensorFlow.jsを使ってjavascript上でkerasモデルを使う - Qiita テンソルをdisposeで解放またはtidy内で行う必要があるとのこと