ユニファ開発者ブログ

ユニファ株式会社プロダクトデベロップメント本部メンバーによるブログです。

SageMakerを用いたオリジナルTensorFlowモデルの推論

こんにちは、データエンジニアリングチームの宮崎です。

少し今更ですが、最近SageMakerにオリジナルのTensorFlowモデルをデプロイできることを知ったので試してみました。 (これまでSageMakerで学習したモデルしかデプロイできないと思っていました…。)

モチベーション

以前、コチラの記事でAWS Lambda上にTensorFlowモデルのデプロイについてご紹介しました。 しかし、AWS Lambdaですとコールドスタート時は応答に20秒前後かかり、応答速度が求められるタスクには辛いところがあります。 また、常に多数のクライアントから推論リクエストがくるケースにも不向きです。 そういったケースはTensorFlow Servingが良さそうですが、ECSにデプロイしようとすると、ロードバランサなどの設定も必要で少々面倒です。

そこで今回は、マネージドでTensorFlow ServingをデプロイしてくれるSageMakerを用い、オリジナルのTensorFlowモデルをデプロイしてみました。

Saved Modelの出力

はじめにSaved Modelを出力します。 Saved Modelの作成にはコチラの記事model_exporter.py を使用します。

$ python model_exporter.py --export_path imagenet-resnet50/1

カスタムハンドラの作成

次に、カスタムハンドラを作成します。 SageMakerのデフォルトハンドラはRESTで推論リクエストを送信するため、画像データをBase64エンコードする必要があります。 そしてTensorFlow Serving側は画像データを受け取るため、Saved ModelにBase64のデコード処理を予め組み込んでおく必要があります。

今回は既にモデルを作成済みのため、Saved Modelをそのまま利用できるgRPCのリクエスト送信をカスタムハンドラに実装したいと思います。 また、合わせてカスタムハンドラにはS3から画像を取得する処理も組み込みます。 これによって、クライアント側はS3キーを送るだけで、よしなに推論してくれるエンドポイントを作ることができます。 したがって全体の構成としては以下のようになります。

f:id:unifa_tech:20210719164109p:plain
SageMakerの構成図

カスタムハンドラのコードは以下のようになります。 注意点として、カスタムハンドラの関数名は handler とする必要があります。 また、context.grpc_port にgRPCポートが設定されているため、これを使って送信先アドレスを作成します。 preprocess で入力されたS3キーから画像をダウンロードし、TensorFlow ServingにgRPCでリクエストを送信します。 そして受信したレスポンスを postprocess でJSONに変換し、クライアントに返します。

def handler(serialized_data, context) -> Tuple[str, str]:
    data = json.loads(serialized_data.read().decode('utf-8'))
    images = preprocess(data)

    server = f'localhost:{context.grpc_port}'
    with grpc.insecure_channel(server) as channel:
        stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
        request = predict_pb2.PredictRequest()
        request.model_spec.name = 'imagenet-resnet50'
        request.model_spec.signature_name = 'serving_default'
        images_tensor_proto = make_tensor_proto(images)
        request.inputs['images'].CopyFrom(images_tensor_proto)
        response = stub.Predict(request, 20.0)

    return postprocess(response)


def preprocess(data: dict) -> np.ndarray:
    s3_object = s3_client.get_object(
        Bucket=data['bucket'], Key=data['s3_key'])
    image_data = io.BytesIO(s3_object['Body'].read())
    image = Image.open(image_data).convert('RGB')
    images = np.expand_dims(np.array(image), axis=0)

    return images


def postprocess(data: predict_pb2.PredictResponse) -> Tuple[str, str]:
    result = {}
    for key, value in data.outputs.items():
        if value.dtype == types_pb2.DT_FLOAT:
            val = np.array(value.float_val)
        elif value.dtype == types_pb2.DT_STRING:
            val = np.array(value.string_val).astype(str)
        else:
            continue

        result[key] = val.reshape(
            [dim.size for dim in value.tensor_shape.dim]).tolist()

    result['model_version'] = int(data.model_spec.version.value)
    response_content_type = 'application/json'

    return json.dumps(result), response_content_type

gRPCの利用にあたっては事前にgRPCライブラリをビルドしておきます。 TensorFlowパッケージを入れても良いのですが、エンドポイントデプロイ時の時間を短縮するため、自前でビルドしました。 また、ここで最新バージョンのgrpcio-toolsでビルドしてしまうと、SageMaker側とバージョンが合わずエラーになってしまうので、注意が必要です。

$ pip3 --no-cache-dir install grpcio==1.29.0 grpcio-tools==1.29.0
$ git clone --depth 1 -b v2.3.1 https://github.com/tensorflow/tensorflow.git
$ cd tensorflow
$ python3 -m grpc_tools.protoc \
    -I. \
    --python_out=../code \
    --grpc_python_out=../code \
    tensorflow/core/{framework,example,protobuf}/*.proto

モデルファイルの作成およびアップロード

最終的にSageMakerにデプロイするモデルファイルは以下のようになります。 カスタムハンドラのコードは code ディレクトリ に格納する必要があります。

.
├── code
│   ├── inference.py    # カスタムハンドラを記述したコード
│   ├── requirements.txt
│   └── tensorflow    # gRPCライブラリ
│       └── core
│           ├── example
│           │   ├── example_parser_configuration_pb2.py
│           │   ├── ...
│           │   └── feature_pb2_grpc.py
│           ├── framework
│           │   ├── allocation_description_pb2.py
│           │   ├── ...
│           │   └── versions_pb2_grpc.py
│           └── protobuf
│               ├── autotuning_pb2.py
│               ├── ...
│               └── worker_service_pb2_grpc.py
└── imagenet-resnet50    # Saved Model
    └── 1
        ├── assets
        ├── saved_model.pb
        └── variables
            ├── variables.data-00000-of-00001
            └── variables.index

Saved Modelおよびカスタムハンドラのコードを圧縮し、モデルファイルとしてS3にアップロードします。

$ tar cvfz model.tar.gz imagenet-resnet50 code
$ aws s3 cp model.tar.gz s3://<bucket_name>/<dir_name>/

SageMakerエンドポイントのデプロイ

アップロードしたモデルファイルをSageMaker推論エンドポイントにデプロイします。 今回はSageMaker Python SDKを使用しました。

from sagemaker.tensorflow import TensorFlowModel

model = TensorFlowModel(
    model_data='s3://<bucket_name>/<dir_name>/model.tar.gz',
    entry_point='inference.py',
    role='arn:aws:iam::<account>:role/<sagemaker_exec_role>',
    framework_version='2.3.1',
    name='imagenet-resnet50')

model.deploy(
    initial_instance_count=1,
    instance_type='ml.t2.medium',
    endpoint_name='imagenet-resnet50')

推論リクエストの送信

デプロイが成功したら、推論リクエストを送信してみます。 今回もこちらの画像を使用しました。

f:id:unifa_tech:20210519093104j:plain
推論に使用した画像 (出典: Open Images Dataset)

import boto3

sagemaker_client = boto3.Session().client('sagemaker-runtime')

response = sagemaker_client.invoke_endpoint(
    EndpointName='imagenet-resnet50',
    Body=json.dumps({
        'bucket': '<bucket_name>',
        's3_key': 'elephant.jpg'
    }),
    ContentType='application/json'
)

result = json.load(response['Body'])
for label, confidence in zip(
        result['labels'][0], result['confidences'][0]):
    print(f'{label:<20}: {confidence:0.6f}')

推論結果です。ちゃんとインド象と推論してくれました。

Indian_elephant     : 0.888790
tusker              : 0.099783
African_elephant    : 0.011377
water_buffalo       : 0.000025
hippopotamus        : 0.000006

まとめ

SageMakerにオリジナルのTensorFlowモデルをデプロイしました。 SageMakerはECSより若干インスタンス料金が増えますが、 ロードバランサを意識しなくてよくなるので便利です。 また、カスタムハンドラによって、前処理・後処理も柔軟に行うことができました。

一方で、カスタムハンドラ周りについては資料が少なく、ディレクトリ名や関数名が決まっているなど、少しとっつきづらかったです。 前後の処理を知りたい方はこちらのコードを参照すると、理解が深まるかと思います。 いずれにしても、無事にデプロイできたので、この先活用できればと思っています。

ソースコード

今回使用したソースコードです。

github.com

参考資料


ユニファで一緒に働く仲間を募集しています!

unifa-e.com