こんにちは。メルカリアドベントカレンダー 9日目は JP AI Engieering Team 所属の @KosukeArase がお送りします。
機械学習技術をサービスに適用するにあたって、機械学習モデルによる推論を API として提供（サービング）したいという場面は多いと思います。
TensorFlow community は TensorFlow Serving によるモデルのサービングを可能にしており、メルカリでも多くの機械学習プロダクトでこの TensorFlow Serving を利用しています。
TensorFlow Serving はドキュメントがしっかりしているため本稿では詳細に説明しませんが、実運用に際してはなるべくモデルサイズを小さくしたい、無駄な計算を減らし推論速度をあげたい、といったモチベーションがあるかと思います。
このような計算グラフの最適化を実現する関数、ツールはいくつかありますが、メルカリでは写真検索などの機能で Graph Transform Tool を利用しています。この Graph Transform Tool は多機能かつ有用なのですが、特に日本語での文献は現状ほとんど見当たらないため、本稿では Graph Transform Tool により計算グラフの最適化を行い、モデルのサイズやサービング環境におけるパフォーマンスの比較を行います。
実験は以下の環境で行いました。
tensorflow == '1.15.0' tensorflow_hub == '0.7.0'
事前準備
例として、Open Images V4 で学習済みの MobileNet V2 を特徴抽出器とした SSD をサービング用に最適化することを考えます。
今回は TensorFlow Hub を利用しグラフの構築と重みのロードを行います。SavedModel や checkpoint などその他の入力形式を利用する場合にも、同様に tf.Graph() にグラフをロードすればそれ以降の処理は同じです。
モデルの詳細はこちら に記載されていますが、入力はサイズ (1, height, width, 3) の tf.float32 である必要があるため、placeholder で画像を受け取り変換したものをグラフへの入力とします。
また、出力である net は dictionary になっていますが、これらは reshape して output という別の dictionary に入れておきます。
import tensorflow as tf import tensorflow_hub as hub from tensorflow.tools.graph_transforms import TransformGraph graph = tf.Graph() with graph.as_default(): module_url = "https://tfhub.dev/google/openimages_v4/ssd/mobilenet_v2/1" input_image = tf.placeholder(dtype=tf.uint8, name="input_image", shape=[1, None, None, 3]) images = tf.squeeze(input_image, axis=[0]) images = tf.image.convert_image_dtype(images, dtype=tf.float32) images = tf.expand_dims(images, 0) model = hub.Module(module_url) net = model(images, as_dict=True) inputs = { "inputs": input_image, } detection_boxes = tf.reshape(net["detection_boxes"], (1, -1, 4)) detection_classes = tf.reshape(net["detection_class_labels"], (1, -1)) detection_scores = tf.reshape(net["detection_scores"], (1, -1)) outputs = { "detection_boxes": detection_boxes, "detection_classes": detection_classes, "detection_scores": detection_scores, }
TensorFlow Serving でモデルをサービングするためには、SavedModel 形式で保存する必要があります。まずは上で読み込んだモデルをそのまま保存してみます。
output_path = "./models/1" with tf.Session(graph=graph) as sess: tf.saved_model.simple_save(sess, output_path, inputs=inputs, outputs=outputs)
出力された SavedModel を確認してみます。
$ ll ./models/1 -rw-r--r-- 1 kosuke.arase staff 122M 12 9 17:42 saved_model.pb drwxr-xr-x 2 kosuke.arase staff 64B 12 9 17:42 variables/
saved_model.pb と言う 122 MB のファイルが保存されていることがわかります。また、variables ディレクトリの中は空です。
Graph Transform Tool による変換
TensorFlow のモデルは推論時には不必要なノードも多く含んでいるため、モデルのサイズが必要以上に大きく無駄な計算をしてしまいます。
例えば、metrics に関連する tf.summary など学習時にのみ必要なノードを取り除いたり、常に定数を評価するノードを事前に定数で置き換えたり、batch normalization で掛け算される係数を事前に Conv2D や MatMul の重みに掛けておくことでまとめたりすることができます。
これらの変換を行うためのツールとして Graph Transform Tool があります。ドキュメント内では Bazel で build し CLI として利用していますが、今回は Python API を使って変換してみます。
使い方は非常に簡単で、GraphDef 形式のグラフ、入出力、施したい変換のリストを tensorflow.tools.graph_transforms.TransformGraph に渡すだけです。
transforms = [ "remove_attribute(attribute_name=_class)", "remove_nodes(op=Identity, op=CheckNumerics)", "strip_unused_nodes", "fold_constants(ignore_errors=true)", "fold_batch_norms", "fold_old_batch_norms", "merge_duplicate_nodes" ] graph_def = graph.as_graph_def() optimized_graph_def = TransformGraph( graph_def, ["input_image:0"], ["Reshape:0", "Reshape_1:0", "Reshape_2:0"], # detection_boxes, detection_classes, detection_scores transforms)
このようにすると、以下のようにそれぞれの変換を適用してくれます。
2019-12-09 18:02:46.982615: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying remove_attribute 2019-12-09 18:02:47.371648: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying remove_nodes 2019-12-09 18:02:51.219189: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying strip_unused_nodes 2019-12-09 18:02:51.578994: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying fold_constants 2019-12-09 18:02:52.848059: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying fold_batch_norms 2019-12-09 18:02:53.294608: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying fold_old_batch_norms 2019-12-09 18:02:55.276499: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying merge_duplicate_nodes
GraphDef を tf.Graph() に戻し、同じように SavedModel として保存してモデルのサイズを見てみましょう。
optimized_graph = tf.Graph() with optimized_graph.as_default(): tf.import_graph_def(optimized_graph_def, name="") output_path = "./models/2" with tf.Session(graph=optimized_graph) as sess: tf.saved_model.simple_save(sess, output_path, inputs=inputs, outputs=outputs)
$ ll ./models/2 -rw-r--r-- 1 kosuke.arase staff 58M 12 9 18:16 saved_model.pb drwxr-xr-x 2 kosuke.arase staff 64B 12 9 18:16 variables/
saved_model.pb のサイズが 122 MB から 58 MB へと小さくなっていることがわかります。
TensorFlow Serving によるサービング
最適化されたモデルを TensorFlow Serving によってサービングしてみましょう。
公式から tensorflow/serving と言う名前で docker image が提供されているので、それを利用します。
$ docker pull tensorflow/serving
$ docker run -p 8500:8500 --rm -v /path/to/models:/models/ssd_mobilenet_v2 -e MODEL_NAME=ssd_mobilenet_v2 -t tensorflow/serving
2019-12-09 09:45:07.032648: I tensorflow_serving/model_servers/server.cc:85] Building single TensorFlow model file config: model_name: ssd_mobilenet_v2 model_base_path: /models/ssd_mobilenet_v2
2019-12-09 09:45:07.035438: I tensorflow_serving/model_servers/server_core.cc:462] Adding/updating models.
2019-12-09 09:45:07.035662: I tensorflow_serving/model_servers/server_core.cc:573] (Re-)adding model: ssd_mobilenet_v2
2019-12-09 09:45:07.144277: I tensorflow_serving/core/basic_manager.cc:739] Successfully reserved resources to load servable {name: ssd_mobilenet_v2 version: 10}
2019-12-09 09:45:07.144341: I tensorflow_serving/core/loader_harness.cc:66] Approving load for servable version {name: ssd_mobilenet_v2 version: 10}
2019-12-09 09:45:07.144368: I tensorflow_serving/core/loader_harness.cc:74] Loading servable version {name: ssd_mobilenet_v2 version: 10}
2019-12-09 09:45:07.144911: I external/org_tensorflow/tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: /models/ssd_mobilenet_v2/10
2019-12-09 09:45:08.017806: I external/org_tensorflow/tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { serve }
2019-12-09 09:45:08.330064: I external/org_tensorflow/tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-12-09 09:45:09.020575: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:202] Restoring SavedModel bundle.
2019-12-09 09:45:09.022285: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:212] The specified SavedModel has no variables; no checkpoints were restored. File does not exist: /models/ssd_mobilenet_v2/10/variables/variables.index
2019-12-09 09:45:09.022362: I external/org_tensorflow/tensorflow/cc/saved_model/loader.cc:311] SavedModel load for tags { serve }; Status: success. Took 1877445 microseconds.
2019-12-09 09:45:09.093024: I tensorflow_serving/servables/tensorflow/saved_model_warmup.cc:105] No warmup data file found at /models/ssd_mobilenet_v2/10/assets.extra/tf_serving_warmup_requests
2019-12-09 09:45:09.103661: I tensorflow_serving/core/loader_harness.cc:87] Successfully loaded servable version {name: ssd_mobilenet_v2 version: 10}
2019-12-09 09:45:09.109854: I tensorflow_serving/model_servers/server.cc:353] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2019-12-09 09:45:09.114076: I tensorflow_serving/model_servers/server.cc:373] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 238] NET_LOG: Entering the event loop ...
リクエストを送る gRPC の Python クライアントは例えばこのような感じで書けます。
image = np.expand_dims(np.array(Image.open("image_path")), 0) request = predict_pb2.PredictRequest() request.model_spec.name = 'ssd_mobilenet_v2' request.model_spec.signature_name = "serving_default" request.inputs["inputs"].CopyFrom(tf.make_tensor_proto(image, dtype=tf.uint8)) with grpc.insecure_channel("localhost:8500") as channel: stub = prediction_service_pb2_grpc.PredictionServiceStub(channel) response = stub.Predict(request)
まとめ
TensorFlow Hub からモデルをダウンロードし、Graph Transform Tool でサービング用にモデルを最適化し、TensorFlow Serving によりサービングするまでの流れをざっくりと追いました。
今回のモデルですと推論速度はほぼ変わりませんが、モデルや含まれているノードによっては上の最適化により推論速度も変わることがあるので、TensorFlow のモデルをサービングする際に上記の最適化はやっておいて損はないものだと思います。
注意点としては、TF v2 で諸々の API が変更されている点です。本当は TF v2 で書きたかったのですが、時間がなく v1 に甘んじてしまいました。すみません。
明日の執筆担当は、メルカリ iOS チームの@hideさんです。それでは引き続きお楽しみください！
