この素晴らしい TensorFlow Serving に祝福を! - docker-compose で簡単デプロイ

はじめに

TensorFlow 2.0 から TensorFlow Serving が実装され，Flask で推論サーバを構築しなくても Docker を使い簡単にエンドポイントを作成できるようになりました．

今回は TensorFlow 2.0 と Docker と docker-compose を使い，簡単にエンドポイントを立てる方法を備忘録としてまとめました．

学習

今回は fashion_mnist を使い簡単にモデルを作成しました．

import json

import tensorflow as tf
from tensorflow import keras

fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
train_images = train_images / 255.0

d = {'signature_name': 'serving_default',
     'inputs': [test_images[0].tolist()]}

with open('./test_data.json', mode='w') as f:
    f.write(json.dumps(d))
    
model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28), name='inputs'),
    keras.layers.Dense(128, activation=tf.nn.relu),
    keras.layers.Dense(10, activation=tf.nn.softmax)
    ])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=5)

tf.saved_model.save(model, './models/fashion_model/1')

エンドポイントを立てる

上記で保存したモデルを使いエンドポイントを立てます．

イメージは公式が公開している tensorflow/serving を使用します．
https://hub.docker.com/r/tensorflow/serving

version: '3'
services:
  app:
    image: tensorflow/serving 
    container_name: fashon_endpoint_test
    ports:
      - 8501:8501
    volumes:
      - ./models:/models/fashion_model
    environment:
      - MODEL_NAME=fashion_model
      - MODEL_BASE_PATH=/models/fashion_model

簡単デプロイ！

$ docker-compose up

テスト

$ curl -X POST -H "Content-Type: application/json" http://localhost:8501/v1/models/fashion_model:predict -d @test_data.json

楽勝

おわりに

昨今では PyTorch がディープラーニングフレームワークのデファクトスタンダードになりつつあり，kaggle や 学術分野などでも多くのコードが PyTorch で書かれています．

PyTorch にはこのような機能がなく(少なくとも自分が調べた範囲では…もしあったら…ごめん)，モデルデプロイに関しては TensorFlow に軍配が上がった感覚を覚えました．

今回は触れませんでしたが，signature やメタデータの利用により，より可用性のあるモデル管理ができるらしいのですが，今回はデプロイから推論のお手軽さのみの紹介にとどめます．

コード

参考