TensorFlow Serving の簡単な負荷試験 - ユニファ開発者ブログ

はじめに

こんにちは、iOSエンジニアのしだです。開発ブログは久々な気がします。今回は TensorFlow Serving のgRPCの簡単な負荷試験してみます。
ECS などで TensorFlow Serving を運用する際に、どれくらいのマシンリソースだとどれくらい処理できるということを確認するためにやってます。

TensorFlow Serving はTensorFlowの学習済みのモデルをAPI化したいときに便利です。 Flaskなどで独自にWebサーバーを用意しなくてもSavedModelを用意するだけでAPIとして提供できます。負荷試験ツールは ghz というフレームワークを利用します。

TensorFlow Serving のDockerイメージが提供されているので、それと学習済みの物体検出モデルを使ってAPIを用意します。そして ghz を使って gRPC の負荷試験できるようにローカルで動かすまでやりたいと思います。

www.tensorflow.org

準備

macOS 10.14.6
Python 3.7.4
TensorFlow 1.14.0

Python ライブラリ

今回必要なライブラリをインストールしておきます。

$ pip install grpcio-tools \
Pillow \
tensorflow==1.14.0 \
tensorflow-serving-api==1.14.0 \
numpy==1.16.4

ghz

github.com

リポジトリのリリースにあるファイルをダウンロードして、適当な場所に解凍して使うことにします。

$ curl -L -O https://github.com/bojand/ghz/releases/download/v0.41.0/ghz_0.41.0_Darwin_x86_64.tar.gz
$ tar zxf ghz_0.41.0_Darwin_x86_64.tar.gz --directory ~/bin
$ ghz -v
0.41.0

物体検出のコンテナイメージの作成

Tensorflow detection model zoo にあるモデルを使って、TensorFlow Serving で物体検出のAPIを提供してみたいと思います。

Dockerイメージの作成手順は、TensorFlow Serving with Dockerを参考に行います。使うモデルは Tensorflow detection model zoo にある ssdlite_mobilenet_v2_coco_2018_05_09.tar.gz を利用させてもらいます。

// モデルファイルのダウンロード
$ curl -L -O http://download.tensorflow.org/models/object_detection/ssdlite_mobilenet_v2_coco_2018_05_09.tar.gz
$ tar zxf ssdlite_mobilenet_v2_coco_2018_05_09.tar.gz
$ mkdir object_detection

// SavedModelのディレクトリ名を適当なバージョン名にします
$ mv ssdlite_mobilenet_v2_coco_2018_05_09/saved_model object_detection/1234

// Dockerイメージにモデルファイルをコピーします
$ docker pull tensorflow/serving:1.14.0
$ docker run -d --name serving_base tensorflow/serving:1.14.0
$ docker cp object_detection serving_base:/models/object_detection
$ docker commit --change "ENV MODEL_NAME object_detection" serving_base object_detection_serving
$ docker kill serving_base
$ docker rm serving_base

TensorFlow Servingを動かしてみて、以下のようなレスポンスが返ってくれば動作しています。

// TensorFlow Serving の起動
$ docker run -p 8500:8500 -p 8501:8501 --rm -it object_detection_serving

// 別ターミナル等で TensorFlow Servingの動作確認
$ curl http://localhost:8501/v1/models/object_detection
{
 "model_version_status": [
  {
   "version": "1234",
   "state": "AVAILABLE",
   "status": {
    "error_code": "OK",
    "error_message": ""
   }
  }
 ]
}

proto ファイルをかき集める

ghz を実行するときにメッセージのprotoファイルが必要になります。 tensorflow/serving にあるprotoファイルは、tensorflow/tensorflow のprotoに依存しているのでそれぞれ取ってきます。（もう少し良いやり方がありそうな気がします。。） tensorflow/tensorflow と tensorflow/serving それぞれのリポジトリをクローンしてきます。

$ git clone --depth 1 -b v1.14.0 https://github.com/tensorflow/tensorflow.git
$ git clone --depth 1 -b 1.14.0 https://github.com/tensorflow/serving.git

先に、tensorflow/core/framework に含まれているprotoをPython用にビルドします。バイナリデータを作成するときに、Pythonから画像データをTensorにしているので必要になります。

$ mkdir proto
$ cd tensorflow
$ python -m grpc_tools.protoc -I . --python_out=../proto --grpc_python_out=../proto tensorflow/core/framework/*.proto
$ cd ..

また、tensorflow_serving/apis にある protoファイルが読み込めるように、tensorflow/core の下にある protoファイルを tensorflow_serving/apis の下にコピーしてくる必要があります。以下のようなディレクトリ構成になるようにファイルをコピーします。

$ mkdir -p tensorflow_serving/apis/tensorflow/core/framework
$ mkdir -p tensorflow_serving/apis/tensorflow/core/example
$ mkdir -p tensorflow_serving/apis/tensorflow/core/protobuf
$ cp serving/tensorflow_serving/apis/*.proto tensorflow_serving/apis
$ cp tensorflow/tensorflow/core/framework/*.proto  tensorflow_serving/apis/tensorflow/core/framework
$ cp tensorflow/tensorflow/core/example/*.proto  tensorflow_serving/apis/tensorflow/core/example
$ cp tensorflow/tensorflow/core/protobuf/*.proto  tensorflow_serving/apis/tensorflow/core/protobuf
$ tree tensorflow_serving
tensorflow_serving
└── apis
    ├── classification.proto
    ├── ...
    ├── session_service.proto
    └── tensorflow
        └── core
            ├── example
            │   ├── example.proto
            │   ├── example_parser_configuration.proto
            │   └── feature.proto
            ├── framework
            │   ├── allocation_description.proto
            │   ├── ...
            │   └── versions.proto
            └── protobuf
                ├── autotuning.proto
                ├── ...
                └── worker_service.proto

バイナリデータの作成

ghz ではバイナリファイルも扱うことができるので、最後に TensorFlow Serving へのリクエストをシリアライズして書き出します。ここでは、画像1枚を predict のメッセージとしてシリアライズして書き出します。

import os
import sys
import argparse

import numpy as np

from tensorflow_serving.apis import predict_pb2

from PIL import Image

from proto.tensorflow.core.framework import tensor_pb2, tensor_shape_pb2, types_pb2
Dim = tensor_shape_pb2.TensorShapeProto.Dim


def main(args):
  if not args.image:
    raise ValueError('You need to --image option.')

  # 画像を読み込む
  image = Image.open(args.image).convert('RGB')
  image.thumbnail((args.image_size, args.image_size))
  image = np.array(image)

  # tensorflow_serving のリクエストを作成します
  request = predict_pb2.PredictRequest()
  request.model_spec.name = 'object_detection'
  request.model_spec.signature_name = args.signature_name

  image_np = np.expand_dims(image, axis=0)
  print('Image:', image_np.shape, image.dtype)

  # リクエストに読み込んだ画像データをセットする
  dim = [Dim(size=i) for i in image_np.shape]
  shape = tensor_shape_pb2.TensorShapeProto(dim=dim)
  content = image_np.tostring()
  tensor = tensor_pb2.TensorProto(
    dtype=types_pb2.DT_UINT8, tensor_shape=shape, tensor_content=content)
  request.inputs['inputs'].CopyFrom(tensor)

  os.makedirs(args.output_dir, exist_ok=True)
  name = os.path.splitext(os.path.basename(args.image))[0]
  filename = '{}-{}.bin'.format(name, str(args.image_size))
  output_path = os.path.join(args.output_dir, filename)

  # シリアライズしたファイルを書き出す
  with open(output_path, 'wb') as f:
    f.write(request.SerializeToString())


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument(
    '--image', type=str, help='path to image in JPEG format')
  parser.add_argument(
    '--image_size', type=int, help='image size', default=800)
  parser.add_argument(
    '--output_dir', type=str, help='output directory', default='/tmp/')
  parser.add_argument(
    '--signature_name',
    type=str, help='signature name', default='serving_default')
  main(parser.parse_args(sys.argv[1:]))

上記のようなシリアライズするコードを用意して、画像ファイルを指定して実行します。実行すると HU020_350A-800.bin というファイル名のバイナリデータが書き出されます。このファイルが負荷試験するときにメッセージになります。

$ python write_bin.py --image /tmp/350A/HU020_350A.jpg --image_size 800 --output_dir ./

実行

最後に ghz を使って実行してみます。 --proto オプションに tensorflow/serving のpredictionのprotoファイルを指定して、 -B オプションに先程作ったシリアライズしたファイルを指定します。

$ ghz --insecure \
--proto tensorflow_serving/apis/prediction_service.proto \
--call tensorflow.serving.PredictionService.Predict \
-B ./HU020_350A-800.bin \
-c 10 \
-n 100 \
localhost:8500

Summary:
  Count:    100
  Total:    3.73 s
  Slowest:  651.90 ms
  Fastest:  160.53 ms
  Average:  357.17 ms
  Requests/sec: 26.82

Response time histogram:
  160.528 [1]    |∎∎
  209.665 [2]    |∎∎∎
  258.803 [11]   |∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎
  307.940 [26]   |∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎
  357.077 [15]   |∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎
  406.215 [17]   |∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎
  455.352 [14]   |∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎∎
  504.490 [4]    |∎∎∎∎∎∎
  553.627 [3]    |∎∎∎∎∎
  602.764 [4]    |∎∎∎∎∎∎
  651.902 [3]    |∎∎∎∎∎

Latency distribution:
  10% in 245.79 ms
  25% in 278.54 ms
  50% in 351.30 ms
  75% in 425.01 ms
  90% in 523.00 ms
  95% in 570.93 ms
  99% in 651.90 ms

Status code distribution:
  [OK]   100 responses

上記はローカルで実行している結果ですが、TensorFlow Servingを使って平均 357.17 ms になります。今回使ったモデルは、Tensorflow detection model zoo によると、画像サイズ600px、TITAN Xで 27 ms で処理するようなのですが、なかなか同じ環境は用意して実行することはできませんが一つ指標にはなります。

最後に

TensorFlow Serving が便利なので最近はよく利用してます。 gRPCの理解も充分ではないですが、protoファイルを理解しながら今回試すことができました。
あとは ECSなどで動かしてマシンリソースを上げれば処理数が上がるとか、要件に合わせてCPUでよいとかGPU使う必要があるとか、確認しながらうまく使っていけたらと思います。

参考

https://ghz.sh/docs/usage