Lobeで学習したモデルとTensorFlow Servingを使ってREST APIを作る。

下記の記事でLobeでの学習と学習したmodelをエクスポートしてPythonから利用する方法を説明しました。

マイクロソフトが公開した機械学習モデルの訓練を容易にできる「Lobe」を試してみた。
Lobeで学習したモデルをPythonで利用する方法

今回はTensorFlow Servingと組み合わせてREST APIを作ってみたいと思います。

TensorFlow Servingを動作させる下準備

Lobeでの学習とmodelのエクスポートは上に記載した記事をご覧ください。

いろいろと試したところTensorFlow Servingを動作させるにはmodelが入っているディレクトリの名前が数字である必要があるようです。
（この辺りはまだはっきりと掴めていないのでご存知の方がいれば教えていただけるとありがたいです。）

そのため、Lobeからエクスポートしたmodelディレクトリの中に数字を名前にしたディレクトリを作成し必要なファイルをコピーもしくは移動します。

DockerでTensorFlow Servingを立ち上げる

TensorFlow Servingを使用する一番簡単な方法はDockerを使う方法です。
docker hubに公式のImageがアップされているのでこれを使用します。
tensorflow/serving

今回使用するmodelのディレクトリは下記のようになっています。
TensorFlow Servingでは１というディレクトリにあるmodelを使用します。

── sample_model
    ├── 1 # 新たに作成しsample_model配下の必要なファイルをコピーもしくは移動する
    │   ├── saved_model.pb
    │   └── variables
    │       ├── variables.data-00000-of-00001
    │       └── variables.index
    ├── example
    │   ├── README.md
    │   ├── requirements.txt
    │   └── tf_example.py
    ├── saved_model.pb
    ├── signature.json
    └── variables
        ├── variables.data-00000-of-00001
        └── variables.index

dockerの起動は下記のように行います。

# docker run -t --rm -p 8501:8501 -v "`pwd`/sample_model:/models/sample_model" -e MODEL_NAME=sample_model tensorflow/serving

modelが読み込まれ正常に動作がスタートすると下記のメッセージが表示されます。

2020-11-03 01:04:51.142942: I tensorflow_serving/model_servers/server.cc:387] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 238] NET_LOG: Entering the event loop ...

このメッセージが出力されていればTensorFlow Servingの起動は完了です。

RESTで画像分類を試す

ここまでの手順でRESTで画像分類を試す下準備は整いました。
TensorFlow Servingの動作を確認するためGETリクエストを送ってみます。

import requests
import json

url = 'http://localhost:8501/v1/models/sample_model'
res = requests.get(url)
print(json.loads(res.text))

下記のようなレスポンスが返ってくれば問題なく動作しています。

{'model_version_status': [{'version': '1', 'state': 'AVAILABLE', 'status': {'error_code': 'OK', 'error_message': ''}}]}

予測のリクエストはPOSTで行います。

import json

import cv2
import numpy as np
import requests

# 予測リクエストURL
url = 'http://localhost:8501/v1/models/sample_model:predict'

def predict(image_path):
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    image = image.astype(np.float32) / 255
    image = image.tolist()

    headers = {"content-type": "application/json"}
    body = {"inputs": [image]}
    r = requests.post(url, data=json.dumps(body), headers=headers)
    r_json = json.loads(r.text)
    return r_json['outputs']['Prediction'][0]

predictに画像ファイルを引数として与えて予測結果を得ることができます。

predict_label = predict('画像ファイル')
print(predict_label)

ここまでで最低限動く環境は作ることができます。

コードはgitにアップしてありますのでよければご覧ください。
lobe_py

FastAPIと組み合わせてREST APIを作ってみました。