Amazon SageMaker で複数の TensorFlow モデルを一つのエンドポイントへデプロイする方法

概要

Amazon SageMakerでは、TensorFlow、MXNet、Chainer、PyTorch、scikit-learnといった機械学習フレームワークをサポートしています。これらのフレームワークを利用して機械学習による予測結果を得るときは、学習した機械学習モデルをエンドポイントにデプロイすることができます。複数のモデルが必要になる場合、これらのモデル一つ一つに対してエンドポイントを作成する方法が一般的ですが、推論リクエストが少ないモデルに対してエンドポイントを常時起動すると、推論処理に対するコストが高くなってしまいます。そこで、推論リクエストの少ないモデルを、他のモデルと同じエンドポイントにデプロイし、常時起動するインスタンス数を低減する方法があります。この手法はリアルタイム処理が必要な場合は特に有効です。なお、推論処理がリアルタイム性を要求しない場合はバッチ変換ジョブをご利用ください。

本記事では、複数のモデルを一つのエンドポイントにデプロイする方法について説明いたします。この方法は、Amazon SageMaker がサポートする全ての機械学習フレームワークで実装可能ですが、ここでは Tensorflow Serving を利用して、複数のモデルをデプロイする方法について説明します。例として、軽量な物体検出モデル SSD MobileNet と、軽量な画像分類モデル MobileNet を1つのインスタンスにデプロイします。デプロイまでの手順の概要は以下のとおりです。

複数のTensorFlowモデルを TensorFlow Serving にデプロイ可能な SavedModel形式で保存します。
保存したモデルを1つのアーカイブファイル (tar.gz 形式) にして、 Amazon S3 (S3) にアップロードします。
Amazon SageMaker のAPIを利用して、1つのインスタンスにデプロイし、テストします。

それでは各手順について以下で説明します。

1. TensorFlowモデルの保存

Jupyter Notebookからコードを実行し、TensorFlowの学習済みモデルをダウンロードして、以下のようなSavedModel形式で保存します。model1をSSD MobileNet、model2をMobileNetとします。TensorFlowモデルには、SavedModel形式以外の学習済みモデルが公開されている場合があるので、必要に応じて変換します。今回はSSD MobileNetとMobileNetの両モデルを変換します。

./multi
├── model1
│ └── 0
│ ├── saved_model.pb
│ └── variables
│ ├── variables.data-00000-of-00001
│ └── variables.index
│
└── model2
└── 0
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index

1-1. 物体検出用モデルの取得

TensorFlow detection model zoo から、SSD MobileNet の学習済みモデルをダウンロード・解凍し、SavedModel形式に変換します。解凍したファイルの中から、学習済みモデルのファイルfrozen_inference_graph.pbをTensorFlowのグラフとして読み込み、builder.save() を利用して保存します。ルート(この場合は /multi/model1/0/ )を指定すれば、上記の通りに保存されます。また、デプロイした際に、 bounding box や score なども返すように signature を定義します。

export_dir = 'multi/model1/0/'
graph_pb = 'ssd_mobilenet_v1_coco_2018_01_28/frozen_inference_graph.pb'
        
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

with tf.gfile.GFile(graph_pb, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    
sigs = {}

with tf.Session(graph=tf.Graph()) as sess:
    # name="" is important to ensure we don't get spurious prefixing
    tf.import_graph_def(graph_def, name="")
    g = tf.get_default_graph()
    input_image = g.get_tensor_by_name("image_tensor:0")
    detection_boxes = g.get_tensor_by_name('detection_boxes:0')
    detection_scores = g.get_tensor_by_name('detection_scores:0')
    detection_classes = g.get_tensor_by_name('detection_classes:0')
    num_detections = g.get_tensor_by_name('num_detections:0')
    
    sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
        tf.saved_model.signature_def_utils.predict_signature_def(
            {"input": input_image},
            {"boxes": detection_boxes,
             "scores": detection_scores,
             "classes": detection_classes,
             "n_detect": num_detections} )

    builder.add_meta_graph_and_variables(sess,
                                         [tag_constants.SERVING],
                                         signature_def_map=sigs)

    builder.save()

1.2. 画像分類用モデルの取得

画像認識用のモデル MobileNet の学習済みモデルはこちらからダウンロードできます。こちらはTF-Slim用のモデル形式なので、モデルのファイルmobilenet_v1_0.25_128.ckptを読み込んで、SavedModel形式に保存します。モデルを読み込むため、学習済みモデルとともに提供されているPythonスクリプトをダウンロードし、 import mobilenet_v1 として利用します。先ほどと同様に、 multi 以下にmodel2/0 を作成して、モデルを保存します。variables 以下のファイルも保存する必要があります。

checkpoint_file = 'mobilenet_v1_0.25_128.ckpt'
# Initialize the network.
tf.reset_default_graph()

# Define the network and load checkpoint
input_tensor = tf.placeholder(tf.float32, shape=(None,128,128,3), name='input_image')
sess = tf.Session()

import mobilenet_v1
with tf.contrib.slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope()):
    logits, end_points = mobilenet_v1.mobilenet_v1(inputs=input_tensor, depth_multiplier=0.25, is_training=False, num_classes = 1001)
saver = tf.train.Saver()
saver.restore(sess, checkpoint_file)

export_path = "./multi"

tf.saved_model.simple_save(
    sess,
    os.path.join('multi/model2/0/'),
    inputs={'input_image': input_tensor},
    outputs={'output': end_points['Predictions']})

2. 二つのモデルを一つのファイルにまとめて、S3 にアップロード

フォルダ./multi 以下に2種類のモデルが保存されました。これを1つのアーカイブファイル multi.tar.gz にして S3 にアップロードします。アップロード先は変数 models に格納され、デプロイするときに参照されます。

!tar -czvf multi.tar.gz ./multi

import sagemaker
from sagemaker import get_execution_role
sagemaker_session = sagemaker.Session()
models = sagemaker_session.upload_data(path='./multi.tar.gz', key_prefix='pretrained_model/multiple_models/')

3. 複数モデルをデプロイし、テストする

アップロードした multi.tar.gz を利用して複数のモデルをデプロイします。デフォルトでどちらのモデルを推論に使用するか、環境変数 env で設定しておきます。ここでは物体検出用モデルがデフォルトで使用されるようにしています。deploy()を呼び出してデプロイを行う方法は、一つのモデルをデプロイする通常の場合と同じです。

from sagemaker.tensorflow.serving import Model, Predictor

role = get_execution_role()

env = {
  'SAGEMAKER_TFS_DEFAULT_MODEL_NAME': 'model1'
}

model = Model(model_data=models, role=role, framework_version='1.12', env=env)
predictor = model.deploy(initial_instance_count=1, instance_type='ml.c5.xlarge')

推論のテストを行うために次の犬の画像を利用します。こちらの画像をJupyter Notebookと同じディレクトリにdog.jpgというファイル名で保存してください。

あらかじめ物体検出と画像分類それぞれのラベルを1行ずつ記述したcoco_labelmap.txtとmobilenet_labelmap.txtを用意します。これらもJupyter Notebookや画像ファイルと同じディレクトリに保存してください。ここでは、これらのテキストファイルを読み込んで、ラベルIDに対応するラベルの名前がわかるようにします。

# Load labelmap for object detection
object_label_map = []
with open('coco_labelmap.txt', 'r') as f:
    for line in f:
        object_label_map.append(line.strip())

#Load labelmap for image classification
cls_label_map = []
with open('mobilenet_labelmap.txt', 'r') as f:
    for line in f:
        cls_label_map.append(line.strip())

それでは実際に SSD MobileNet を使って、物体検出の推論を行ってみます。

%matplotlib inline
from PIL import Image
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np


# Load image and resize it for inference.
pil_img = Image.open('dog.jpg')
height, width, _ = np.array(pil_img).shape
#Preprocess for base model
image = pil_img.resize((128,128)).convert('RGB')
image = np.array(image)
im =  image.reshape(128,128, 3)

#Inference by base model
result = predictor.predict({"input":im})

# Extract results of which score is larger than 0.5
scores = result['predictions'][0]['scores']
good_scores = [score for score in scores if score > 0.5]

# Parse result data
rect = []
label = []
score = []
for i in range(len(good_scores)):
    # Extract box data and generate rectangle objects to be draw on image later
    box = result['predictions'][0]['boxes'][i]
    y_min = box[0]*height
    x_min = box[1]*width
    y_max = box[2]*height
    x_max = box[3]*width
    im_width = x_max - x_min
    im_height = y_max - y_min 
    rect.append(patches.Rectangle((x_min,y_min),im_width,im_height,linewidth=1,edgecolor='b',facecolor='none'))
    
    # Store the label name
    label_id = result['predictions'][0]['classes'][i]
    label.append(object_label_map[int(label_id)])
    
    # Store score
    score.append(good_scores[i])

# Display the image with bounding boxes, labels and scores.
fig,ax = plt.subplots(1)
ax.imshow(np.array(pil_img))
for index in range(len(rect)):
    ax.add_patch(rect[index])
    caption = label[index] + ": " + str(score[index])
    plt.text(x_min,y_min-10,caption,fontsize='12', color = 'b')
    print(caption)
plt.show()

こちらのような出力が得られるはずです。

次に画像分類を試します。画像分類モデルはデフォルトに設定されたモデルではないため、画像分類モデル model2 を呼び出すPredictorを作成します。predict を呼び出すことによって、今度は画像分類の結果を呼び出すことができます。

%matplotlib inline
# get the endpoint name from the default predictor
endpoint = predictor.endpoint

# Load image and resize it for inference.
pil_img = Image.open('dog.jpg')
height, width, _ = np.array(pil_img).shape
#Preprocess for base model
image = pil_img.resize((128,128)).convert('RGB')
# image = (np.array(image)/255 - 0.5 )*2
image = np.array(image)

im =  np.array(image.reshape(-1, 128,128,3))
# get a predictor for 'model2'
model2_predictor = Predictor(endpoint, model_name='model2')

class_id=np.argmax(model2_predictor.predict(im)['predictions'][0]['probabilities'])
probability = model2_predictor.predict(im)['predictions'][0]['probabilities'][class_id]

print(cls_label_map[class_id] + ": " + str(probability))
fig,ax = plt.subplots(1)
ax.imshow(np.array(pil_img))
plt.show()

出力のサンプルは次のとおりです。

これで2つのモデルが1つのエンドポイントへとデプロイできることが確認できました。最後に、動作確認が完了し、必要のなくなったエンドポイントはsagemaker.Session().delete_endpoint(predictor.endpoint) で削除しましょう。

まとめ

本稿ではTensorFlowのデプロイ環境としてTensorFlow Servingを利用し、複数モデルをAmazon SageMakerで作成した一つのエンドポイントへとデプロイする方法を示しました。複数モデルを一つのエンドポイントへとデプロイすることでコストを抑えながらリアルタイムに推論する環境を提供することができます。Amazon SageMakerでは今回ご紹介した機械学習モデルのデプロイ以外にも、機械学習モデルの開発やトレーニング、データのラベル付け、ハイパーパラメータ最適化、トレーニング済モデルのコンパイルなどの機能がお使い頂けます。皆様がAmazon　SageMakerを利用して機械学習を使ったサービスを構築する際は、このブログ記事がご参考になれば幸いです。

著者について

okunotom

奥野友哉 (Tomoya Okuno, Ph.D.) は AWS のソリューションアーキテクトで、反応流体計算・宇宙・機械学習の分野に興味があります。休日はコーヒーが美味しいカフェで過ごしています。

samejim

鮫島正樹 (Masaki Samejima, Ph.D.) は AWS の機械学習スペシャリストソリューションアーキテクトで、コンピュータビジョンや時系列解析などを得意にしています。

Amazon Web Services ブログ