TensorFlow ServingでTensorFlowのモデルを運用する

2019-02-11

TensorFlow ServingとはTensorFlowのモデルをマネージしてくれるサーバです。モデルを処理するサーバ機能に加えて、複数モデルの管理、バージョンの管理、新たなモデルの自動ロード等を提供してくれます。通信はgRPCやRESTで行うため、言語を問わずこれを立ち上げるだけで機械学習の恩恵を受けられる優れものです。 TensorFlow公式で、Google社内でも使っているっぽいのである程度の安心感があります。

モデルの用意

まずTensorFlowのモデルを用意します。今回のチュートリアルにあるFashion MNISTを使います。
https://www.tensorflow.org/tutorials/keras/basic_classification

Fashion MNISTは28*28のサイズに洋服の絵が描かれたデータセットです。 10種類のカテゴリからなり、一般的によく使われる手書きのデータセット(MNIST)よりもデータのパターンが多く、新たなスタンダードを狙っているっぽいです。

作り方はチュートリアルの通りにpythonで作成します。だいたいこんな感じのコードでモデルが作れます。 pb形式でモデルを保存する必要があるため、simple_saveを使って保存します。

また、テスト用のデータを以下のスクリプトで作っておきます。あとでrubyからも利用するので２つ作り、正しい結果も確認しておきます。それぞれtest_data.jsonは9、ruby_data.jsonは2が正解です。

TensorFlow Servingの起動

公式がDockerを用意しているのでそれを利用します。
https://hub.docker.com/r/tensorflow/serving

起動は以下の通りで、./modelsフォルダに先程作ったモデルを置いておきます。

docker run -p 8500:8500 -p 8501:8501 -v `pwd`/models:/work -e MODEL_NAME=fmnist_model -e MODEL_BASE_PATH=/work -t tensorflow/serving

TensorFlow ServingはREST APIを提供しているので、先程作ったデータを使って試してみます。以下のエンドポイントで実行できます。
http://{host}:{port}/v1/models/{model_name}:predict

curl -X POST -H "Content-Type: application/json" http://localhost:8501/v1/models/fmnist_model:predict -d @test_data.json
{
    "outputs": [
        [
            5.28375e-06,
            3.96955e-07,
            6.83376e-07,
            4.72291e-07,
            2.60042e-06,
            0.0151633,
            1.0217e-05,
            0.0781524,
            0.00012393,
            0.906541
        ]
    ]
}%

最も高いのはcategory 9で、正解データと一致します。これでTenserflow Servingの準備はおしまいです。

rubyから繋ぐ

Tenserflow Servingをプログラムから利用します。POSTに送ることで結果を得られますが、gRPCも提供しているためそちらを利用してみます。

定義ファイルの生成

TenserflowとTenserflow Servingにあるprotoファイルから対応する.rbファイルを作ります。なお、servingフォルダの直下にTenserFlowをcloneしてください。なお、2019/02/10現在、Ruby 2.6ではgoogle-protobufが動かないので注意が必要です。(google-protobufの3.7.0で治る)

git clone --depth 1 -b 1.12.0 [email protected]:tensorflow/serving.git
git clone --depth 1 -b v1.12.0 [email protected]:tensorflow/tensorflow.git
mv tensorflow/ serving/

# docker run -it -v `pwd`:/work -t grpc_ruby /bin/bash

gem install grpc grpc-tools

grpc-toolsにはprotoファイルをコンバートするのに必要なprotocコマンドと、ruby用のプラグインのgrpc_ruby_pluginが含まれています。なお、gemで入れた場合はprotocコマンドはgrpc_tools_ruby_protocに、grpc_ruby_pluginはgrpc_tools_ruby_protoc_pluginという名前になっているのでご注意ください。

protoファイルから.rbファイルを生成します。ここでは/work以下にcloneを行ったと仮定し、/work/proto_rubyフォルダ以下に出力しています。

gPRCを利用してTenserFlow Servingに送る

先程作ったruby_data.jsonを使ってTenserFlow Servingにデータを送り、その結果を得ます。やることとしてはgPRCを使って必要なデータを送るだけです。

終わりに

今回は使いませんでしたがTenserFlow Servingでは複数モデルの運用や、複数バージョンの運用ができます。その際も、model_specに指定するデータを変えるだけでアクセスが可能なため、TensorFlowのモデルを運用するために必要な機能は一通りまかなえそうです。