TensorFlow v2 Serving 1分クッキング

- arata-furukawa

TensorFlow(バージョン2以降)を使うと、演算グラフの構築からAPIサーバのデプロイまでたったの1分でできてしまいます。前半では1分でAPIサーバを立てる例を示し、後半で詳細を解説します。

手順はたったのこれだけです:

  1. デプロイしたい演算をTensorFlowで構築する
  2. 演算グラフをSavedModel形式でエクスポートする
  3. SavedModelをTensorFlow ServingのDockerコンテナでデプロイする

なお、この記事では以下の環境で動作確認をしています:

  • Python 3.8.2
  • TensorFlow 2.2.0
  • Docker 19.03.9

1分クッキング

例として入力を2倍するだけの演算グラフ twice(x)=2x\text{twice}(x) = 2xのAPIサーバを1分で作ります。

デプロイしたい演算をTensorFlowで構築する

import tensorflow as tf
print(tf.__version__) # v2以降

class MyMod(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float64)])
    def twice(self, x):
        return 2 * x

これは以下のように実行することができます。

mod = MyMod()
print(mod.twice(2)) # => tf.Tensor(4.0, shape=(), dtype=float64)
print(mod.twice(5)) # => tf.Tensor(10.0, shape=(), dtype=float64)

演算グラフをSavedModel形式でエクスポートする

mod = MyMod()
tf.saved_model.save(mod, 'models/mymod/1')

カレントディレクトリのmodels/mymod/1以下にSavedModel形式で出力されます。必ず<model_base_dir>/<model_name>/<model_version_number>のディレクトリレイアウトで出力してください

以下のようにファイルが出力されていることを確認します。

$ tree models/mymod/1
models/mymod/1
├── assets
├── saved_model.pb
└── variables
    ├── variables.data-00000-of-00001
    └── variables.index

2 directories, 3 files

SavedModelをTensorFlow ServingのDockerコンテナでデプロイする

$ docker run --rm \
  -v $(pwd)/models:/models \
  -p 8501:8501 \
  -e MODEL_NAME=mymod \
  tensorflow/serving \
  --enable_batching

これでlocalhost:8501でREST APIが公開されました!

動作確認

$ curl -sS localhost:8501/v1/models/mymod:predict -XPOST -d '{"inputs":6}'
{
    "outputs": 12.0
}

$ # バージョン指定もできます
$ curl -sS localhost:8501/v1/models/mymod/versions/1:predict -XPOST -d '{"inputs":6}'

TensorFlow Servingの公開するAPI仕様の詳細はこちらのドキュメントを参考にしてください。

解説

TensorFlow Servingを使うと、TensorFlow SavedModel形式の任意の演算グラフを即時にハイパフォーマンスAPIサーバとしてデプロイできます。TensorFlow ServingはC++で書かれたサーバで、リクエストを並列で処理しながら高速に演算グラフを実行します。

配信にはSavedModelというTensorFlowの演算グラフのシリアライズ形式の一つを用います。Modelという命名から、機械学習モデルを想像しますが、実際には任意の演算グラフをシリアライズできます。

演算グラフの構築

前半の例を見てみます。重要なポイントが2つあります。

  1. 新しいtf.Moduleを定義している
  2. tf.functioninput_signatureを指定している
class MyMod(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float64)])
    def twice(self, x):
        return 2 * x

tf.Module

SavedModel形式にエクスポートできるのはTrackableなオブジェクトです。これは、tf.Moduletf.train.Checkpointtf.keras.Modelなどが例として挙げられます。

Trackableなオブジェクトは属性として持っているノードを追跡しています。例えば、以下のように定義しても同じ結果が得られます。

@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float64)])
def twice(x):
    return 2 * x

mod = tf.Module()
mod.twice = twice

実際に機械学習モデルをデプロイする場合は、tf.keras.ModelをそのままSavedModelにすることもできますし、モデルを属性に持ったtf.Moduleを使えばモデルも一緒にエクスポートされます。

class MyModel(tf.Module):
    def __init__(self):
        self.model = tf.keras.applications.MobileNet()

    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float64)])
    def predict(self, x):
        return self.model(x)

mod = MyModel()
tf.saved_model.save(mod, 'models/mymodel/1')
loaded_mod = tf.saved_model.load('models/mymodel/1')
loaded_mod.predict(tf.zeros([1,224,224,3],dtype=tf.float64)).shape
# => TensorShape([1, 1000])

tf.functioninput_signature

v2で追加されたtf.functionを使うことで、簡単に演算グラフを拡張できます。tf.functionはPython関数をTensorFlowの演算グラフに自動変換するので、if文やfor文などのPython構文をそのまま使うことができます。

そして重要な点が、input_signatureで入力のシグネチャ(引数のテンソルの形状、データ型)を明示することです。上記ガイドにも記載があるように、input_signatureがない場合は、関数の呼び出し時に入力シグネチャに合わせて関数の再トレーシングが発生します。そのおかげでポリモーフィックに振る舞うことができるようになっているのですが、TensorFlow Servingでデプロイするためには、エクスポートの段階でSavedModelにシグネチャを具体的に明示する必要があります。

SavedModelシグネチャの指定

前半の例では、以下のようにSavedModelにエクスポートしました。

mod = MyMod()
tf.saved_model.save(mod, 'models/mymod/1')

このとき、TensorFlowはmodの中で明示されたシグネチャを探し、一つだけ見つかった場合、それをserving_defaultというシグネチャ名として自動的に採用します。明示的にシグネチャを指定する場合、以下のようにも書けます。

# すべて同じ結果になります
tf.saved_model.save(mod, 'models/mymod/1')
tf.saved_model.save(mod, 'models/mymod/1', signatures=mod.twice)
tf.saved_model.save(mod, 'models/mymod/1', signatures={
    'serving_default': mod.twice,
})

もしシグネチャが見つからなかった場合、serving_defaultシグネチャがSavedModelに保存されません。SavedModelを単純にモデルの保存だけに使うなら問題ありませんが、TensorFlow Servingではシグネチャの指定がない場合はserving_defaultシグネチャを読み込んで使用しようとするため、APIサーバを構築するなら指定しておくことが推奨されます。

シグネチャの活用

ちなみに、このシグネチャは自由に使うことができ、Servingでserving_default以外のシグネチャを使用することもできますし、複数のシグネチャを保存することもできます。

@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float64)])
def twice_x100(x):
    return 100 * mod.twice(x)

tf.saved_model.save(mod, 'models/mymod/1', signatures={
    'serving_default': mod.twice,
    'my_signature': twice_x100,
})

このようにmodに含まれていない関数でも大丈夫です。追加したシグネチャはパラメータのsignature_nameで指定することで利用できます。

$ curl -sS localhost:8501/v1/models/mymod:predict -XPOST -d '{"inputs":6}'
{
    "outputs": 12.0
}

$ curl -sS localhost:8501/v1/models/mymod:predict -XPOST -d '{"inputs":6, "signature_name":"my_signature"}'
{
    "outputs": 1200.0
}

例えば、モデルがRGB画像のビットマップである3階のテンソル[28,28,3]を受け取るように実装されている場合に、serving_defaultではそのまま3階の[28,28,3] tf.floatテンソルを受け取れるようにしておき、一方でJPEGやPNGでエンコードされたデータNone tf.stringを受け取り画像をデコードしてからモデルに適用する別のシグネチャを用意する、といった使い方ができます。

メタデータの取得

APIで配信されているモデルのバージョンや、利用可能なシグネチャ、入力の形状や型情報などを取得できるmetadataエンドポイントが利用できます。

前半の例のレスポンスを見てみると、serving_defaultシグネチャが定義されているのがわかります。

$ curl -sS localhost:8501/v1/models/mymod/metadata
{
"model_spec":{
 "name": "mymod",
 "signature_name": "",
 "version": "1"
}
,
"metadata": {"signature_def": {
 "signature_def": {
  "serving_default": {
   "inputs": {
    "x": {
     "dtype": "DT_DOUBLE",
     "tensor_shape": {
      "dim": [],
      "unknown_rank": false
     },
     "name": "serving_default_x:0"
    }
   },
   "outputs": {
    "output_0": {
     "dtype": "DT_DOUBLE",
     "tensor_shape": {
      "dim": [],
      "unknown_rank": false
     },
     "name": "PartitionedCall:0"
    }
   },
   "method_name": "tensorflow/serving/predict"
  },
  "__saved_model_init_op": {
   "inputs": {},
   "outputs": {
    "__saved_model_init_op": {
     "dtype": "DT_INVALID",
     "tensor_shape": {
      "dim": [],
      "unknown_rank": true
     },
     "name": "NoOp"
    }
   },
   "method_name": ""
  }
 }
}
}
}

サーバ設定をカスタマイズする

設定ファイルを使用することで、複数モデル・複数バージョンの配信や、公開するバージョンのポリシー(デフォルトで最新バージョンのみ)を変更することなどができます。

モデルの配信設定

例えば、

  • model1model2model3モデルを配信する(例なので同じデータを使っていますがもちろん別のモデルデータを配信できます)
  • model1はバージョン2と3だけを公開する
  • model2は最新バージョンだけを公開する
  • model3は全バージョンを公開する
# model.config
model_config_list {
  config {
    name: 'model1'
    base_path: '/models/mymod'
    model_platform: 'tensorflow'
    model_version_policy {
      specific {
        versions: 2
        versions: 3
      }
    }
  }
  config {
    name: 'model2'
    base_path: '/models/mymod'
    model_platform: 'tensorflow'
    model_version_policy {
      latest {}
    }
  }
  config {
    name: 'model3'
    base_path: '/models/mymod'
    model_platform: 'tensorflow'
    model_version_policy {
      all {}
    }
  }
}

設定ファイルをコンテナにマウントして、起動時に--model_config_file=/model.configのようにパラメータを渡します。

TensorFlow Servingは、<base_path>/<name>/<version>のディレクトリを常に監視しており、ポリシーで対象となるServableが発見・削除されると再起動しなくても自動でホットリローディングを行います。ポリシーが最新バージョンや全バージョンの場合、バージョン指定なしでAPI呼び出しを行った場合は最もバージョンの新しい(数字の大きい)モデルを利用します。

ちなみに昔デフォルトであったモデル設定ファイルのホットリローディング機能はなくなりました(テスト目的でポーリングする機能はありますが本番環境での使用は推奨されません)。ポリシーを更新する場合はファイルを更新してコンテナを再起動するか、gRPC経由でtensorflow.serving.ModelService.HandleReloadConfigRequestModelServerConfigを渡して呼び出すことで再読み込みが行われます。gRPCではダウンタイムがないですが、個人的にはk8sデプロイメントなどを活用してコンテナのロールアウトをすることで更新するのをオススメします。

バージョン名のエイリアスを作れるバージョンラベル機能がありますが、2020/06/27現在、RESTでサポートされていません。gRPCで使用する場合は、stableなどのラベルを使うと開発が便利です。

バッチ処理の設定

シグネチャがバッチ処理に対応している場合、--enable_batchingを使うことでTensorFlow Servingの処理が最適化されます。他にもキューなどの処理の制限や並列処理スレッド数などを制御できます。--batching_parameters_file=/batch.configのように指定します。

# /batch.config
max_batch_size {
  value: 128
}
batch_timeout_micros {
  value: 0
}
max_enqueued_batches {
  value: 1000000
}
num_batch_threads {
  value: 8
}

モニタリング設定

Prometheus形式でメトリクスを取得することができます。--monitoring_config_file=/monitoring.configで指定します。

# /monitoring.config
prometheus_config {
  enable: true,
  path: "/metrics/prometheus"
}

以下のようにメトリクスが公開されます。

$ curl -sS localhost:8501/metrics/prometheus
# TYPE :tensorflow:api:op:using_fake_quantization gauge
# TYPE :tensorflow:cc:saved_model:load_attempt_count counter
:tensorflow:cc:saved_model:load_attempt_count{model_path="/models/mymod/2",status="success"} 1
# TYPE :tensorflow:cc:saved_model:load_latency counter
:tensorflow:cc:saved_model:load_latency{model_path="/models/mymod/2"} 38157
# TYPE :tensorflow:cc:saved_model:load_latency_by_stage histogram
:tensorflow:cc:saved_model:load_latency_by_stage_bucket{model_path="/models/mymod/2",stage="init_graph",le="10"} 0
:tensorflow:cc:saved_model:load_latency_by_stage_bucket{model_path="/models/mymod/2",stage="init_graph",le="18"} 0
... 長いので略 ...
# TYPE :tensorflow:mlir:import_failure_count counter
# TYPE :tensorflow:serving:model_warmup_latency histogram
# TYPE :tensorflow:serving:request_example_count_total counter
# TYPE :tensorflow:serving:request_example_counts histogram
# TYPE :tensorflow:serving:request_log_count counter

まとめ

TensorFlow v2でDefine-by-Runがデフォルトになり、tf.functionのような柔軟なグラフ構築のサポートが追加されたことで、任意のTensorFlowの演算グラフを簡単にSavedModelにエクスポートできるようになり、TensorFlow Servingを使うことですぐにAPIサーバを作れるようになりました。

TensorFlow Servingは十分にテストされ、運用されている実績があります。コンテナですぐに始められるだけでなく、バージョンコントロールやホットリローディング、メトリクスの収集やバッチ処理の制限とチューニングなど、本番環境での運用に最低限必要な機能がすべて備わっています。今回は紹介しませんでしたが、TensorFlow以外のモデルをTensorFlow Servingで配信する方法もあり、拡張性に富んでいます。

v2はv1からの破壊的変更が大きく、未だにv2に移行できていないプロジェクトも多いですが、使いこなせれば強いツールとなります。TensorFlow v2でServingを活用してみてくださいね。