本記事では、tensorflowで作成したモデルをTnesorflow.jsで使用できる形式へ変換する方法を示します。

python環境

pip install tensorflowjs

tfjs-converter

Tensorflow.js(TF.js)は、Tensorflow(TF)によって学習された既存のモデルを再利用可能。tfjs-converterは、TensorFlowモデルを変換するためのコマンドラインツールで、HDF5など様々な形式をサポートしている。

tensorflowjs_converter --help
usage: TensorFlow.js model converters. [-h]
                                       [--input_format {tensorflowjs,keras,tf_hub,keras_saved_model,tf_saved_model,tfjs_layers_model}]
                                      [--output_format {tfjs_graph_model,tfjs_layers_model,tensorflowjs,keras}]
                                       [--signature_name SIGNATURE_NAME]
                                       [--saved_model_tags SAVED_MODEL_TAGS]
                                       [--quantization_bytes {1,2}]
                                       [--split_weights_by_layer] [--version]
                                       [--skip_op_check SKIP_OP_CHECK]
                                       [--strip_debug_ops STRIP_DEBUG_OPS]
                                       [--weight_shard_size_bytes WEIGHT_SHARD_SIZE_BYTES]
                                       [input_path] [output_path]

TFのSavedModelを、TF.jsで読み取れるWeb形式へ以下のコマンドで変換できる。

tensorflowjs_converter \
    --input_format=tf_saved_model \
    /path/to/saved_model \
    /path/to/web_model

--input_formatは, keras, keras_saved_model, tf_hubなど、他の形式も選択可能。

portable形式

モデルが訓練される場所とは異なる場所で再利用できる。特徴は、

軽量: 限られたメモリ要領に格納できる
シリアライズ: ネットワークI/Oを介して共有可能
互換性: 複数のプラットフォームで使用可能

Tensorflowによって生成されるファイルは、プロトコルバッファーに基づいている。そのため、多くのプログラミング言語で読み取り可能な形式です。
また、ONNXなどの形式なども、プラットフォームに依存しない形式です。(pytorchなどでも使用可能)

モデルの保存

Tensorflow(Python)でモデルを保存する方法です。

mport tensorflow as tf
from tensorflow import keras

print(tf.__version__) #2.1

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

train_labels = train_labels[:1000]
test_labels = test_labels[:1000]

train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

# 短いシーケンシャルモデルを返す関数
def create_model():
  model = tf.keras.models.Sequential([
    keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10, activation='softmax')
  ])
  
  model.compile(optimizer='adam', 
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])
  
  return model

# 基本的なモデルのインスタンスを作る
model = create_model()
model.summary()

model.fit(train_images, train_labels, epochs=5)

# モデル全体を１つのHDF5ファイルに保存します。
model.save('my_model.h5')
tf.saved_model.save(model, "./sample/model_data")
#imported = tf.saved_model.load("./sample/model_data")

モデルの変換

./sample/model_dataに保存されたTFモデルをTF.js形式のモデルに変換する。

tensorflowjs_converter --input_format=tf_saved_model --output_node_names=output ./sample/model_data ./sample/model_tfjs_model

#hd5データ形式
tensorflowjs_converter --input_format keras my_model.h5 ./sample/hd5_model

#tf-hubのモデルを使用
tensorflowjs_converter \
    --input_format=tf_hub \
    'https://tfhub.dev/google/imagenet/mobilenet_v1_100_224/classification/1' \
    ./my_tfjs_model

TF-hubとは

Tensorflow Hubは、機械学習のモデルなどを再利用できるライブラリです。
モデルを読み込み、入出力の形式を合わせれば、最先端の機械学習技術を使用できます。

モデルのロード

Tensorflow.jsでモデルロードする方法です。

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

const MODEL_URL = 'https://path/to/model.json';
const model = await tf.loadGraphModel(MODEL_URL);

// Or

const MODEL_PATH = 'file://path/to/model.json';
const model = await tf.loadGraphModel(MODEL_PATH);