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

概要

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

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

1. 複数の TensorFlow モデルを TensorFlow Serving にデプロイ可能な SavedModel 形式で保存します。
2. 保存したモデルを1つのアーカイブファイル (tar.gz 形式) にして、Amazon S3 (S3) にアップロードします。
3. 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() を利用して保存することで SavedModel 形式のモデルを得ることができます。ルート(この場合は ./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:
    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()

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']})

!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/')

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')

# 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())

%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()

%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()