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

昨年、AWS Lambdaにおけるコンテナイメージのサポートが発表され、AWS Lambda上でTensorFlowを使用しやすくなりました。一方で、以前コチラの記事でご紹介したように、AWS LambdaでTensorFlowを使用するにはServerless Python Requirementsを利用する方法もあります。コンテナを使う場合は、ロードなどにより処理時間が増えそうでもあり、今後どのように使い分ければ良いのか迷いどころです。

そこで、今回はServerless Python Requirementとコンテナイメージ利用、それぞれで処理時間がどのようになるのか計測してみました。

使用するソフトウェアのバージョン

今回の計測では以下のソフトウェアを使用しました。特にTensorFlowはバージョンが新しくなるとファイルサイズが大きくなり、Serverless Python Requirementsの手法が使えなくなってしまったため、若干古い2.0.0を使用しました。

ソフトウェア	バージョン	備考
Python	3.7
TensorFlow	2.0.0
Serverless Framework	2.20.1
Serverless Python Requirements	5.1.0	コンテナ未使用時のみ

Serverless Python RequirementsによるTensorFlowの使用

Serverless Python Requirementsを利用した手法については以下にまとめておりますので、ご参照ください。

tech.unifa-e.com

コンテナによるTensorFlowの使用

続いて、コンテナイメージを利用したTensorFlowの使用方法になります。

Saved Modelの出力

初めに、ImageNetの事前学習済みモデルに前処理と後処理を追加して、Saved Modelを出力します。前処理では画像のリサイズとRGBの正規化、後処理では予測値の整形をそれぞれ行います。

import os
import json
import argparse
import tensorflow as tf

MEAN_RGB = [103.939, 116.779, 123.68]
IMAGENET_CLASS_URI = 'https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json'


class ModelExporter:
    def __init__(self, top=5):
        self.top = top
        self.classes = self.load_classes()

    @staticmethod
    def resnet50_preprocess(image, input_shape):
        image = tf.image.resize(image, input_shape)
        image = image[..., ::-1]
        image -= MEAN_RGB
        return image

    @staticmethod
    def load_classes():
        file_path = tf.keras.utils.get_file(
            os.path.basename(IMAGENET_CLASS_URI), IMAGENET_CLASS_URI)
        with open(file_path) as f:
            classes = json.load(f)

        classes = [classes[str(i)][1] for i in range(len(classes))]
        classes = tf.convert_to_tensor(classes)

        return classes

    def build_serve_fn(self, model):
        @tf.function(
            input_signature=[tf.TensorSpec([None, None, None, 3], dtype=tf.uint8)])
        def serve(images):
            images = self.resnet50_preprocess(images, model.input.shape[1:3])
            predictions = model(images)
            confidences, labels = tf.math.top_k(predictions, k=self.top)
            labels = tf.map_fn(
                lambda x: tf.map_fn(
                    lambda y: self.classes[y], x,
                    dtype=tf.string),
                labels, dtype=tf.string)

            return {'labels': labels, 'confidences': confidences}

        return serve

    def export(self, model, export_path):
        serve_fn = self.build_serve_fn(model)
        tf.saved_model.save(model, export_path, signatures={
            'serving_default': serve_fn})


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--export_path', type=str, required=True)
    args = parser.parse_args()

    # Build model
    model = tf.keras.applications.ResNet50(weights='imagenet')

    # Export model
    exporter = ModelExporter()
    exporter.export(model, args.export_path)

$ python model_exporter.py --export_path saved_model

saved_modelディレクトリにSaved Modelが出力されました。

Lambdaハンドラの作成

続いて、出力したSaved Modelを用いて推論を行うLambdaハンドラを作成します。指定されたS3上の画像を読み込み、推論した結果を返します。

import io
import logging
import boto3
import tensorflow as tf

MODEL_PATH = 'saved_model'

s3_client = None
model = None

logger = logging.getLogger()
logger.setLevel(logging.INFO)


def handler(event, _context):
    logger.info(f'Event: {event}')

    global s3_client, model
    if s3_client is None:
        session = boto3.Session()
        s3_client = session.client('s3')

    if model is None:
        model = tf.saved_model.load(MODEL_PATH)

    # Download image
    s3_object = s3_client.get_object(
        Bucket=event['bucket'], Key=event['filename'])
    image = io.BytesIO(s3_object['Body'].read())
    image = tf.io.decode_jpeg(image.getvalue())

    # Load model
    serve = model.signatures['serving_default']
    result = serve(tf.expand_dims(image, axis=0))

    # Format result
    labels = result['labels'][0].numpy()
    confidences = result['confidences'][0].numpy()
    result = [{'class': l.decode(), 'confidence': float(c)}
              for l, c in zip(labels, confidences)]

    logger.info(f'Result: {result}')

    return result

Dockerコンテナの作成

作成したSaved ModelおよびLambdaハンドラを保持するDockerコンテナを作成します。以下のファイルを同一ディレクトリに保存します。

.
├── Dockerfile
├── app.py (Lambdaハンドラ)
├── requirements.txt
├── saved_model (出力したSaved Model)
│   ├── assets
│   ├── saved_model.pb
│   └── variables
│       ├── variables.data-00000-of-00001
│       └── variables.index
└── serverless.yml

Dockerfileには以下のように記述します。Lambda用のPythonコンテナをベースにし、最後にLambdaハンドラの実行を指定します。

FROM public.ecr.aws/lambda/python:3.7

COPY . ${LAMBDA_TASK_ROOT}
RUN pip install -r requirements.txt
CMD ["app.handler"]

requirements.txtは必要なPythonパッケージを記述します。今回はboto3とtensorflowのみ設定しています。

boto3~=1.14.19
tensorflow==2.0.0

準備が完了したら、Dockerコンテナを作成します。

$ docker build -t lambda-imagenet-resnet50 .

ECRにアップロード

Dockerコンテナを作成できたら、AWS Lambdaから呼び出せるようにするため、ECRにアップロードします。事前にECRにコンテナリポジトリを作成しておきます。今回は lambda-imagenet-resnet50 という名前にしました。

続いて、作成したDockerコンテナをアップロードします。

$ aws ecr create-repository \
    --repository-name lambda-imagenet-resnet50

$ aws ecr get-login-password --region <Region> | \
    docker login --username AWS --password-stdin <Account>.dkr.ecr.<Region>.amazonaws.com

$ docker tag lambda-imagenet-resnet50:latest \
    <Account>.dkr.ecr.<Region>.amazonaws.com/lambda-imagenet-resnet50:latest

$ docker push <Account>.dkr.ecr.<Region>.amazonaws.com/lambda-imagenet-resnet50:latest

AWS Lambdaのデプロイ

いよいよLambdaのデプロイです。LambdaのデプロイにはServerless Frameworkを使用します。 serverless.ymlに以下を定義します。Serverless Python Requirementsを使用する場合に比べ、だいぶシンプルに書けます。 S3の画像に対して推論を行うため、S3へのアクセスを許可しています。

service: imagenet-container

provider:
  name: aws
  stage: dev
  region: ap-northeast-1
  iamRoleStatements:
      - Effect: "Allow"
        Action:
          - s3:ListBucket
          - s3:GetObject
        Resource:
          - "arn:aws:s3::*"

functions:
  classify:
    image: ${opt:imageUrl}
    memorySize: 2048
    timeout: 180

serverless.ymlを定義できたら、デプロイします。

$ serverless deploy \
    --imageUrl <Account>.dkr.ecr.<Region>.amazonaws.com/lambda-imagenet-resnet50:latest

推論の実行

デプロイが完了したら、いよいよ推論です。事前に適当なS3に対象の写真をアップロードしておきます。今回も象の写真を使用しました。

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

$ serverless invoke \
    --function classify \
    --data '{"bucket": "<Bucket>", "filename": "elephant.jpg"}'
[
    {
        "class": "Indian_elephant",
        "confidence": 0.8918280005455017
    },
    {
        "class": "tusker",
        "confidence": 0.09677373617887497
    },
    {
        "class": "African_elephant",
        "confidence": 0.011350669898092747
    },
    {
        "class": "water_buffalo",
        "confidence": 0.000023159593183663674
    },
    {
        "class": "hippopotamus",
        "confidence": 0.000006523193405882921
    }
]

どうやらこの像はインド像のようです。

処理時間の比較

さて、それぞれAWS Lambdaをデプロイし、ようやく準備が整いました。それでは処理時間を比較したいと思います。わかりやすさのため、単位は秒に揃えて丸めています。

これまでAWS Lambdaでは状態がコールドスタートとウォームスタートの2段階でしたが、コンテナイメージ利用時はコンテナ更新後が加わりました。各状態について、クライアント側の計測時間およびサーバ側(CloudWatch Logs)のDuration, Init Duration, Billed Durationをまとめています。

理由は不明ですが、コンテナイメージ利用時にコンテナ更新後はAWS Lambdaは2度実行されておりましたのでそれぞれの時間を記載しております。また、コンテナイメージ利用時のコールドスタートでは Init Duration がかかるようになり、 Billed Duration は Duration と Init Duration の合算値となっております。

なお、各計測の試行回数は1回のため、精度にばらつきがあることはご承知おきください。

状態	項目	Serverless Python Requirements	コンテナイメージ
コンテナ更新後	クライアント側	-	78.0 秒
	Duration	-	50.8 秒 + 12.1 秒
	Init Duration	-	0.0 秒 + 4.9 秒
	Billed Duration	-	50.8 秒 + 17.0 秒
コールドスタート	クライアント側	26.3 秒	18.6 秒
	Duration	15.0 秒	12.5 秒
	Init Duration	0.0 秒	5.3 秒
	Billed Duration	15.0 秒	17.8 秒
ウォームスタート	クライアント側	0.4 秒	0.4 秒
	Duration	0.3 秒	0.3 秒
	Init Duration	0.0 秒	0.0 秒
	Billed Duration	0.3 秒	0.3 秒

考察

コンテナイメージ利用時、コンテナ更新後はコンテナロードのため1分強処理に時間がかかるようでした。一方で、コールドスタート時は若干時間が短くなっております。これは、Serverless Python Requirementsを利用する方式において、パッケージの展開やモデルのロード時間がかかるためと思われます。ウォームスタートではほぼ時間は変わりませんでした。

以上の検証から、以下のように使い分けるのが良さそうです。

要件	選択する方式
コンテナのロード時間を許容できる場合	コンテナイメージ
許容できない場合	Serverless Python Requirements

一方で、Serverless Python Requirementsを利用する場合、tmp領域の512MB制限により、最新のTensorFlowの利用ができませんでした。最新のTensorFlowを利用したい場合はAWS LambdaからEFSの利用の導入を検討する必要がありそうです。

今回初めてコンテナイメージの利用を行いました。簡単に様々なPythonパッケージを利用できる点が非常に便利だと感じました。一方で以下の注意点がありましたので合わせて記載します。

タイムアウト時間はコンテナロードの時間を考慮し、長めに設定する
AWS Lambdaからコンテナイメージはハッシュ値で参照されているため、ECR上のイメージを更新したら、AWS Lambda側も更新する必要がある

また、挙動についても分からない点が多いため、引き続き調査が必要そうです。

コンテナロード時になぜ2回実行されるのか
コンテナロードはどれくらいの頻度で起きるのか

まとめ

まとめです。今回はAWS Lambdaにおいて、Serverless Python Requirementsを利用した場合とコンテナイメージを利用した場合、2つのケースにおいてTensorFlowの処理時間を計測しました。検証結果から、コンテナイメージを利用する場合は、コンテナロードに1分ほど余計に時間がかかるため、要件によって使い分けが必要なことがわかりました。

ソースコード

検証に使用したソースコードは以下に格納しました。

github.com

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

unifa-e.com

ユニファ開発者ブログ

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

AWS LambdaのコンテナイメージによるTensorFlow推論の処理時間比較

目次