tfjs-converterでpythonのモデルをtfjsで実行する

Tensorflowのモデル（グラフ）をPython→JSに変換するツールtfjs-converterの簡単なサンプルを紹介する．

Tensorflowが主に使われる深層学習モデルの学習は主にPythonで実装されることが多い．対して学習済みモデルを利用する状況は必ずしもPythonが言語としてベストであるとは限らない．中でもWebブラウザはデモやアプリケーションの実行環境としてよく使われることがある．しかしモデルの推論をPythonに任せる場合にはサーバーを立てる必要がありコストがかかる．軽いモデルの場合にはクライアント側，つまりJavascriptに学習済みのモデルの実行をさせたい．

TensorflowはJavascriptでの実行にも力を入れていて，tfjsを中心に活発に開発されている．今回扱うtfjs-converterはPython版Tensorflowで定義された計算過程(Graph）を変数の値ごとtfjs向けに変換するツールである．これを使うとtfjs側ではGraphの定義を設計者が書くことなく，モデルを丸ごとPythonから移行できる．つまり計算資源の柔軟な操作が得意なPythonで学習をした後にJavascriptで使うときにはインターフェースの開発に集中できる．

サンプルのコードは以下に公開している：

• モデルの保存
• モデルの変換とロード

大まかな流れ

モデルの保存

Tensorflowの公式ガイドに従う．

1. Python版Tensorflowでモデルの入力，出力のテンソルを定義する
2. tf.saved_modelで保存

モデルの変換とロード

1. tfjs-converterでモデルをJavascriptで扱えるように変換
2. tfjsでロードして使う

モデルの保存

まずはモデルを定義しよう：

my_model.py

class MyModel(tf.keras.Model):

    scope = 'my_model'

    def __init__(self):
        super().__init__()

        self.layer1 = tf.keras.layers.Dense(10, activation=tf.nn.tanh)
        self.layer2 = tf.keras.layers.Dense(2, activation=tf.nn.tanh)

        self.input_name = None
        self.output_name = None

    def __call__(self, x):
        """ return output of this model """
        self.input_name = x.name
        with tf.name_scope(self.scope):
            x = self.layer1(x)
            x = self.layer2(x)
            x = tf.identity(x, name="output")
        self.output_name = x.name
        return x

    def get_variables(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.scope)

    def get_io_names(self):
        assert not self.input_name is None and not self.output_name is None, "input_name or output_name is not set"
        return self.input_name, self.output_name

__call__は前向き計算（入力→出力）を行う関数．get_variables, get_io_namesはモデルの変数を返したり入出力のテンソルの名前を取得するだけのお助け関数たち．

今回は変数の最適化はスキップして，乱数で初期化して保存する．

save_model.py

def save_model():
    # get directory for export as pathlib.Path
    export_dir = get_export_dir()

    my_model = MyModel()

    x = tf.placeholder(tf.float32, shape=(1, 2), name="input")
    y = my_model(x)

    with tf.Session() as sess:
        init_model_vars(sess, my_model)

        result = sess.run(y, {x: np.array([[1, 2]], dtype=np.float32)})
        print(f"y = {result}")

        tf.saved_model.simple_save(
            sess,
            str(export_dir),
            inputs={"inputs": x},
            outputs={"outputs": y}
        )

        input_name, output_name = my_model.get_io_names()
        print(f"input name = {input_name}")
        print(f"output name = {output_name}")

my_modelは先ほど定義したモデルのインスタンスになっている．これにTensorであるplaceholderxを入力して出力yを得る．モデルで使っているtf.keras.layers.Denseは前向き計算を実行しないと変数がグラフに追加されないので注意（たぶんtf.get_variableを使っている）．

モデルの保存はtf.Saverではなく，tf.saved_model.simple_saveを使う．ここがポイントで，変数の値だけではなく，計算の過程を表すグラフそのものも保存してくれる．

保存された内容は

saved_model.pb  variables/

のようになっていて， pbファイルにグラフがはいっている．

モデルの変換とロード

tfjs-converterのREADMEにしたがって変換する：

convert.sh

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    ../save-tf-model/tmp/my_model \
    ./tmp/converted_my_model

../save-tf-model/tmp/my_modelはPython版で保存したモデルのディレクトリに対応する．これを./tmp/converted_my_modelに変換する．変換の設定はオプションで渡す．

tfjs-converterを使う上で困ったのは2つ：

1.　Python版Tensorflowと同じpyenv環境にインストールするとPython版の挙動がおかしくなる．今回はpipenvを使ったがtf.Sessionが初期化できなくなった（なぜかは追いかけてない）
2.　tensorflowjs_converterのhelpではoutput_formatは指定する必要が必ずし無いように書いているが，実際には必要

変換が終われば読み込みは非常に簡単:

main.ts

import * as tf from '@tensorflow/tfjs-node';

const MODEL_URL = 'file://./tmp/converted_my_model/model.json';

async function main() {
    const model = await tf.loadGraphModel(MODEL_URL)
    const x: tf.Tensor2D = tf.tensor([[1, 2]])
    const y = await model.predict(x)
    console.log(y.toString())
}

main()

計算過程であるLayerの記述はまったく必要ないことに注目してほしい．Typescriptで書いてあるがjavascriptそのままでも同じことができるはず．

ちなみにPythonでロードする場合にはこうなる：

load_model.py

def load_model():
    export_dir = get_export_dir()

    sess = tf.Session()

    metagraph = tf.saved_model.loader.load(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        export_dir)

    input_name = metagraph.signature_def["serving_default"].inputs["inputs"].name
    output_name = metagraph.signature_def["serving_default"].outputs["outputs"].name

    x = sess.graph.get_tensor_by_name(input_name)
    y = sess.graph.get_tensor_by_name(output_name)

    result = sess.run(y, {x: np.array([[1, 2]], dtype=np.float32)})
    print(f"y = {result}")

入出力のテンソルを明示的にgraphから取得する必要があるのでtfjsと比べてすこし面倒かもしれない．

その他

上記の文中ソースコードは色々省略したので詳細はレポジトリを見てほしい．

tfjs-converterがサポートしているoperation（モデル計算過程の単位）は公式のレポジトリに書いてある．必ずしも任意のモデルが変換可能であるかはわからない．

tf.saved_model.simple_saveがduplicatedになっているのでTensorflow v2から変わるかも．