freeeの開発情報ポータルサイト

トップ > Advent Calendar 2017 > TensorFlow Servingで機械学習モデルをプロダクション環境で運用する

TensorFlow Servingで機械学習モデルをプロダクション環境で運用する

Advent Calendar 2017 技術

こんにちは、freee株式会社でエンジニアをやっている米川（@yonekawa）です。最先端のテクノロジーを使って新しいソリューションを生み出していくことをミッションにした、CTW (Change The World) という役職で働いています。

この記事はfreee Developers Advent Calendar 2017の23日目です。

機械学習ではアルゴリズムや大規模データ処理が注目されがちですが、学習したモデルをどうやってサービスで運用するかも悩ましい問題です。実験やアルゴリズムの検証では強力なツールが揃っているPythonがよく使われるので、そのままPythonでAPI作るケースが多いと思います。しかしプロダクション環境で運用するとなると開発しやすさ以外にも、大量リクエスト時のパフォーマンスやデプロイ、モデルの精度評価やA/Bテストなどさまざまな課題があります。

またfreeeでは、WebサービスはRuby on Rails、バックエンドサーバーはGo + gRPCという構成でマイクロサービス基盤が構築されているため、できればこれらの資産を流用して開発することで生産性や品質を高めたい気持ちがありました。

こういった背景があって、より良い機械学習プロダクトのプロダクション運用ができる方法が無いか考えてみることにしました。 そして、いくつかの選択肢を検討する中でTensorFlow Servingが良さそうだったのでご紹介します。

f:id:yonekawa:20171222150905p:plain
Tensorflow Servingによる機械学習モデルの運用イメージ

TensorFlow Servingとは

TensorFlow ServingはTensorFlowで構築した機械学習モデルをプロダクション環境で運用することを目的に設計されたモジュールです。 以下のような特徴があります。

  • C++で書かれた堅牢で安定した高パフォーマンスなRPCサーバー
  • gRPCを使って機械学習モデルによる推論を呼び出せるインタフェース
  • 新しいモデルを読み込んだり複数のモデルを並行運用したりが簡単にできる
  • モデルの読み込みなど一部の実装がプラガブルになっていて用途に応じて書き換えられる

このように、冒頭で触れた機械学習モデルの運用における課題を解決する機能が一通り揃っていることがわかります。 gRPCで推論リクエストを送れることで言語を問わないところと、モデルのバージョン管理や並行運用がしやすいところがメリットになると思います。 そしてGoogleでも実際に使われていることからパフォーマンスや安定性もある程度は期待できます。（まだベンチマークは取っていないのですが、Googleによると推論の時間やネットワークを除いて100,000クエリ/秒くらいは捌けるらしいです）

TensorFlow Servingを使った機械学習モデルの配信は以下のような流れになります。

  1. 学習済みモデルをエクスポートしてモデルサーバーを起動する
  2. クライアントからgRPCでモデルサーバーに推論リクエストを送る
  3. モデルのバージョンを追加してモデルサーバーに反映する

順を追って解説していきます。

準備: モデルサーバーのインストール

モデルサーバーはソースからビルドすることもできますが、apt-getで簡単にインストールすることができます。

$ echo "deb [arch=amd64] http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-model-server tensorflow-model-server-universal" | sudo tee /etc/apt/sources.list.d/tensorflow-serving.list
$ curl https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | sudo apt-key add -
$ sudo apt-get update && sudo apt-get install tensorflow-model-server

こちらのDockerfileを使うと最小構成でtensorflow-model-serverが動作する環境を準備できます。

$ docker build --pull -t tensorflow-model-server -f Dockerfile .

1. 学習済みモデルをエクスポートしてモデルサーバーを起動する

モデルサーバーで読み込むための学習済みモデルを、SavedModelとしてローカルファイルにエクスポートします。例えばこういうMNISTのモデルがあったとします。

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

sess = tf.InteractiveSession()

images = tf.placeholder('float', shape=[None, 28, 28, 1], name='images')
x = tf.reshape(images, [-1, 784])
y_ = tf.placeholder('float', shape=[None, 10])
w = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

sess.run(tf.global_variables_initializer())
scores = tf.nn.softmax(tf.matmul(x, w) + b, name='scores')
cross_entropy = -tf.reduce_sum(y_ * tf.log(scores))
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
for _ in range(FLAGS.training_iteration):
    batch = mnist.train.next_batch(50)
    train_step.run(feed_dict={x: batch[0], y_: batch[1]})

この学習済みのセッションをtensorflow.saved_model.builder.SavedModelBuilderを使ってSavedModelに変換します。

prediction_signature = (
    tf.saved_model.signature_def_utils.build_signature_def(
        inputs={'images': tf.saved_model.utils.build_tensor_info(images)},
        outputs={'scores': tf.saved_model.utils.build_tensor_info(scores)},
        method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

version = 1
export_path = 'models/mnist/{}'.format(str(version))
builder = tf.saved_model.builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(
    sess, [tf.saved_model.tag_constants.SERVING],
    signature_def_map={
        'predict_images': prediction_signature
    },
    legacy_init_op=tf.group(tf.tables_initializer(), name='legacy_init_op'))

builder.save(as_text=False)

コードの全体はこちらです。モデルが出力できたらモデルサーバーを起動します。起動時のオプションでモデルの名前と配信するディレクトリを指定します。

$ ls models/mnist
1
$ tensorflow_model_server --port=9000 --model_name=mnist --model_base_path=models/mnist

Kerasを使っている場合

TensorFlow 1.4においてTensorFlowがKerasを統合し、KerasのモデルをEstimatorに変換する機能が追加されました。Estimatorはexport_savedmodelでSavedModelとして保存できるので、KarasでもTensorFlow Servingによるモデル配信を活用できます。

import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.estimator.export import export

model = keras.applications.vgg16.VGG16(weights='imagenet')
model.compile(optimizer=keras.optimizers.SGD(lr=.01, momentum=.9),
              loss='binary_crossentropy',
              metrics=['accuracy'])
estimator = tf.keras.estimator.model_to_estimator(keras_model=model)
feature_spec = {'input_1': model.input}
serving_input_fn = export.build_raw_serving_input_receiver_fn(feature_spec)
estimator.export_savedmodel(export_path_base, serving_input_fn)

しかし、こういったバグがあったようで記事を書いている時点で最新のTensorFlowではまだうまく動作しませんでした。 上記Pull Requestの内容をパッチで当てるとうまく動作することは確認できたのでアップデートが待たれるところです。

2. gRPCを使ってモデルサーバーに推論リクエストを送る

TensorFlow ServingはgRPCでエンドポイントを提供しています。Protocol Bufferのインタフェースに従えばGoやRailsから推論を利用できます（もちろん画像データを数値化するなどの前処理が必要になったりはするのですが...）

TensorFlowのリポジトリに配置されているProtocol Bufferのファイルからクライアントコードを自動生成できます。 例えばGoのクライアントコードを生成するには以下のようにします。

$ git clone --recursive https://github.com/tensorflow/serving.git
$ protoc -I=serving -I serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow_serving/apis/*.proto
$ protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/framework/*.proto
$ protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/protobuf/{saver,meta_graph}.proto
$ protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/example/*.proto

$GOPATH/src以下にtensorflowとtensorflow_serving用の自動生成コードが出力されます。 以下のようなコードでgRPCでPredictionServiceを呼び出すことで、モデルサーバーにリクエストを送ることができます。

import (
    protobuf "github.com/golang/protobuf/ptypes/wrappers"
    pb "tensorflow_serving/apis"
    "google.golang.org/grpc"
)

img, err := os.Open("path/to/image")
if err != nil {
    panic(err)
}
defer img.Close()
p, err := png.Decode(img)
if err != nil {
    panic(err)
}
inputTensorValues := make([]float32, 28*28)
for i := 0; i < 28; i++ {
    for j := 0; j < 28; j++ {
        r, _, _, _ := p.At(i, j).RGBA()
        inputTensorValues[i+(j*28)] = float32(r) / 255
    }
}

request := &pb.PredictRequest{
    ModelSpec: &pb.ModelSpec{
        Name:          "mnist",
        SignatureName: "predict_images",
    },
    Inputs: map[string]*tfcoreframework.TensorProto{
        "images": &tfcoreframework.TensorProto{
            Dtype: tfcoreframework.DataType_DT_FLOAT,
            TensorShape: &tfcoreframework.TensorShapeProto{
                Dim: []*tfcoreframework.TensorShapeProto_Dim{
                    &tfcoreframework.TensorShapeProto_Dim{
                        Size: int64(1),
                    },
                    &tfcoreframework.TensorShapeProto_Dim{
                        Size: int64(28),
                    },
                    &tfcoreframework.TensorShapeProto_Dim{
                        Size: int64(28),
                    },
                    &tfcoreframework.TensorShapeProto_Dim{
                        Size: int64(1),
                    },
                },
            },
            FloatVal: inputTensorValues,
        },
    },
}
conn, err := grpc.Dial(*servingAddress, grpc.WithInsecure())
if err != nil {
    panic(err)
}
defer conn.Close()

client := pb.NewPredictionServiceClient(conn)
resp, err := client.Predict(context.Background(), request)
if err != nil {
    panic(err)
}

for i, s := range resp.Outputs["scores"].FloatVal {
    if s == 1.0 {
        fmt.Print(i)
        break
    }
}

コードの全体はこちらです。注意点として上記コードにある通り、TensorShapeProtoは1次元の値しか取れなくなっているため、入力が多次元の場合には1次元に変換する必要があります。本来の構造の情報はDimで指定することでPredictionServiceが正しい形式で取り扱ってくれます。

実際に実行してみると以下のような結果が返ってくると思います（テスト用にmyleott/mnist_pngを使わせていただきました）。ちゃんと推論が動いていますね！

$ go run serving_client/mnist/client.go --serving-address localhost:9000 mnist_png/testing/2/1174.png
2

3. モデルのバージョンを追加してモデルサーバーに反映する

データの増加やアルゴリズムの見直しによって精度が向上した最新の学習モデルに更新したい、というユースケースはよくあります。 TensorFlow Servingのモデルサーバーはローカルファイルシステムを監視しており、新しいモデルが配置されたら自動で読み込んでくれます。 なので --model_base_path に指定されたディレクトリにバージョン番号のディレクトリを作ってモデルを配置するだけで、自動で新しいモデルを使えるようになります。

例として以下のように、学習回数を極端に少なくしたバージョン1を作ってみます。

$ python export/tensorflow_mnist --model_version=1 --training_iteration=1 /tmp/models/mnist
$ ls /tmp/modesl/mnist
1
$ tensorflow_model_server --port=9000 --model_name=mnist --model_base_path=models/mnist
$ go run serving_client/mnist/client.go --serving-address localhost:9000 mnist_png/testing/2/1174.png
8

学習が足りていないため精度が悪く間違った推測をしています。では次に1000回学習させたモデルをバージョン2として配置してみます。サーバーを再起動する必要はありません。

$ python export/tensorflow_mnist --model_version=2 --training_iteration=1000 /tmp/models/mnist
$ ls /tmp/modesl/mnist
1  2
$ go run serving_client/mnist/client.go --serving-address localhost:9000 mnist_png/testing/2/1174.png
2

このようにモデルを配置するだけで起動中のモデルサーバーが自動で最新のモデルを読み込んでくれました。

モデルのバージョンポリシーを変更する

モデルサーバーはデフォルトだと指定されたディレクトリから最新のモデルだけを読み込むため、新しいモデルを配置すると古いモデルは配信されなくなります。そうではなく新旧2つのモデルを並行運用したいとか、入力の仕様を変えた2つのモデルをエンドポイントを変えて運用したいなどの時は、モデルのバージョンポリシーを変更することで対応できます。

バージョンポリシーには以下の3つがあります。

  • Latest: 常に最新のモデルを配信する（デフォルト）
  • All: すべてのモデルを配信する
  • Specific: 特定のバージョンのモデルのみ配信する

モデルサーバーのバージョンポリシーを変更するには設定ファイルを作り、--model_config_fileで読み込む必要があります。 例えばディレクトリ以下にある全てのモデルを配信し続けたい場合は以下のような設定ファイルを作ります（中身はテキスト形式のProtocol Bufferです）。

model_config_list: {
  config: {
    name: "mnist",
    base_path: "/tmp/models/mnist",
    model_platform: "tensorflow",
    model_version_policy: { all: {} },
  },
}
$ ls /tmp/models/mnist
1  2
$ tensorflow_model_server --port=9000 --model_config_file=./misc/model.conf

PredictRequestのModelSpecでVersionを指定することで、任意のバージョンを指定して呼び出すことができます。

import (
    protobuf "github.com/golang/protobuf/ptypes/wrappers"
    pb "tensorflow_serving/apis"
)

request := &pb.PredictRequest{
    ModelSpec: &pb.ModelSpec{
        Name:          "mnist",
        SignatureName: "predict_images",
        Version: &protobuf.Int64Value{ Value: int64(1) },
    },
    Inputs: map[string]*tfcoreframework.TensorProto{
        ....
    },
}

この設定ファイルは入力の仕様を変えた2つのモデルをエンドポイントを変えて運用したい場合にも使えます。configフィールドは繰り返し指定ができるので、名前を変えて2つ定義すれば別のモデルとして呼び出せます。

model_config_list: {
  config: {
    name: "mnist",
    base_path: "/tmp/models/mnist",
    model_platform: "tensorflow",
  },
  config: {
    name: "mnist2",
    base_path: "/tmp/models/mnist2",
    model_platform: "tensorflow",
  },
}

ModelSpecでNameを指定すると、呼び出すモデルをリクエスト側で変更できます。

import (
    protobuf "github.com/golang/protobuf/ptypes/wrappers"
    pb "tensorflow_serving/apis"
)

request := &pb.PredictRequest{
    ModelSpec: &pb.ModelSpec{
        Name:          "mnist2",
        SignatureName: "predict_images"
    },
    Inputs: map[string]*tfcoreframework.TensorProto{
        ....
    },
}

このように3つのバージョンポリシーと設定ファイルを駆使すれば、プロダクション環境における学習モデルの更新やデプロイのユースケースの多くは網羅できるのではないかと思います。

まとめ

今回使ったサンプルコードなどはすべてこちらのリポジトリに置いてありますのでご自由にお使いください。

TensorFlow Servingが嬉しいのは、開発や検証には便利なライブラリを活用できるPythonを使い、プロダクション環境ではパフォーマンスや運用コストを考慮してシステム設計するというフローが手軽に実現できるところです。gRPCで通信できるのでマイクロサービスの構成要素として自然に組み込めますし、もちろんKubernetesに載せることも可能です。学習モデルの更新は要件に依存しますが、バージョンポリシーによってさまざまなユースケースに対応できます。分類や回帰などの一般的なタスクの範囲を超えるモデルの場合は（C++を書くことで）プラガブルにカスタマイズすることも可能です。

もっと楽をしたい場合はフルマネージドなクラウドサービスを検討してもいいと思います。 Azure ML StudioはGUIで簡単に学習処理を記述できますし、Amazon SageMakerならJupyter Notebookで学習処理を書いたモデルをシームレスにプロダクションに公開することができます。（GCPは詳しくないですがCloud Machine Learningというのがあるようです）

機械学習の技術は日進月歩で進化していますが、それをプロダクトに活用するエンジニアリングはまだまだ試行錯誤が必要だと感じています。 freeeでは、機械学習を前提にしたプロダクトの基盤を作りユーザーに高速に価値を届けていく腕力のあるエンジニアを募集しています。 一緒に世界を変えるプロダクトを創りましょう。

www.wantedly.com

jobs.freee.co.jp

明日はいよいよ弊社CTOであるよこじ氏が満を持して登場します。ご期待ください。