Amazon SageMaker エンドポイントとAWS Lambdaを使って、YOLOv5の推論をスケールさせる

この記事は “Scale YOLOv5 inference with Amazon SageMaker endpoints and AWS Lambda” を翻訳したものです。

データサイエンティストが要件を満たせる機械学習 (ML) モデルを考案できたら、組織の他のメンバーが推論に簡単にアクセスできるようにモデルを展開する必要があります。しかし、コストおよび計算効率の両面で最適化された状態で大規模にデプロイすることは、困難で面倒な作業になる可能性があります。Amazon SageMaker エンドポイントは、モデルのデプロイのために簡易な方法で、スケーラブルかつコスト最適化されたソリューションを提供します。YOLOv5 モデルはGPLv3 ライセンスの下で配布され、実行時の効率と検出精度の高さで知られる人気のあるオブジェクト検出モデルです。この投稿では、事前に学習済みの YOLOv5 モデルをSageMaker エンドポイントでホストし、AWS Lambda 関数を使用してこれらのエンドポイントを呼び出す方法を示します。

ソリューション概要

以下の画像は YOLOv5 モデルを SageMaker エンドポイントでホスティングし、Lambda を使ってそのエンドポイントを呼び出すためのAWSサービスの構成です。SageMaker ノートブックでは、Amazon Simple Storage Service (Amazon S3) のバケットにある PyTorch 形式の YOLOv5 モデルにアクセスし、それを TensorFlow の SavedModel 形式のYOLOv5モデルに変換してS3 バケットに保存し直します。このモデルは、エンドポイントをホストするときに使用されます。Amazon S3 に画像をアップロードすると、Lambda 関数をトリガーします。この関数はOpenCVの Lambda レイヤーを使用して、アップロードされたイメージを読み取り、エンドポイントを使用して推論を実行します。推論の実行後、必要に応じて推論から得られた結果を使用できます。

この投稿では、まず PyTorch 形式の YOLOv5 のデフォルトモデルを利用し、それを TensorFlow の SavedModel 形式に変換するプロセスについて説明します。変換されたモデルは、SageMaker エンドポイントを使用してホストされます。その後、エンドポイントを呼び出して推論を実行する Lambda 関数を作成して発行します。事前に学習済みの YOLOv5 モデルは GitHub で入手できます。この投稿では yolov5l モデルを使用することとします。

前提条件

前提条件として、以下のように AWS Identity and Access Management (IAM) ロールを SageMaker, Lambda, Amazon S3 への的確なポリシーを設定する必要があります。

SageMaker IAM ロール – S3へのモデルの格納やアクセスのために AmazonS3FullAccess ポリシーのアタッチが必要になります。
Lambda IAM ロール – このロールには複数のポリシーが必要になります。
- S3に格納された画像へのアクセスのために、以下のポリシーが必要になります。
  - s3:GetObject
  - s3:ListBucket
- SageMaker エンドポイントを実行するために、以下のポリシーが必要になります。
  - sagemaker:ListEndpoints
  - sagemaker:DescribeEndpoint
  - sagemaker:InvokeEndpoint
  - sagemaker:InvokeEndpointAsync

以下のリソースやサービスも必要となります。

Lambdaの作成と設定に、AWS コマンドラインインターフェイス (AWS CLI) が必要になります。
SageMaker ノートブックインスタンスが必要になります。これらにはDockerが事前インストールされてあり、これを使用してLambda レイヤーを作成します。ノートブックインスタンスを設定するために、以下の手順を完了します。
- SageMaker コンソールにおいて、ノートブックインスタンスを作成し、ノートブックに名前、インスタンスタイプ（この記事では ml.c5.large を用います）、IAM ロール、その他のパラメーターを設定します。
- こちらの公開リポジトリをクローンし、Ultralyticsによって提供されている YOLOv5 リポジトリを追加します。

YOLOv5 を SageMaker エンドポイントにホストする

事前に学習済みの YOLOv5 モデルを SageMaker でホストする前に、model.tar.gz の中身をエクスポートし正しい構造にパッケージする必要があります。この投稿ではYOLOv5 を SavedModel 形式でホストする方法を示します。YOLOv5 リポジトリには、さまざまな方法でモデルをエクスポートできる export.py ファイルがあります。YOLOv5リポジトリをのクローンし、コマンドラインから YOLOv5 ディレクトリに入ると、以下のコマンドでモデルをエクスポートできます。

$ cd yolov5
$ pip install -r requirements.txt tensorflow-cpu
$ python export.py --weights yolov5l.pt --include saved_model --nms

このコマンドは yolov5 ディレクトリの中に yolov5l_saved_model という新しいディレクトリを作成します。 yolov5l_saved_model ディレクトリは以下のようになっています。

yolov5l_saved_model
    ├─ assets
    ├─ variables
    │    ├── variables.data-00000-of-00001
    │    └── variables.index
    └── saved_model.pb

model.tar.gz を作成するために、yolov5l_saved_model の中身を export/Servo/1 に移動します。コマンドラインから以下のコマンドを実行して export ディレクトリを圧縮し、S3 バケットにアップロードします。

$ mkdir export && mkdir export/Servo
$ mv yolov5l_saved_model export/Servo/1
$ tar -czvf model.tar.gz export/
$ aws s3 cp model.tar.gz "<s3://BUCKET/PATH/model.tar.gz>"

その後、以下のコードを利用してSageMakerノートブックからSageMaker エンドポイントにデプロイすることができます。

import os
import tensorflow as tf
from tensorflow.keras import backend
from sagemaker.tensorflow import TensorFlowModel

model_data = '<s3://BUCKET/PATH/model.tar.gz>'
role = '<IAM ROLE>'

model = TensorFlowModel(model_data=model_data,
                        framework_version='2.8', role=role)

INSTANCE_TYPE = 'ml.m5.xlarge'
ENDPOINT_NAME = 'yolov5l-demo'

predictor = model.deploy(initial_instance_count=1,
                            instance_type=INSTANCE_TYPE,
                            endpoint_name=ENDPOINT_NAME)

上記のスクリプトは、モデルをSageMaker エンドポイントにアップロードするのに約2-3分かかります。デプロイのステータスはSageMakerのコンソール上で確認することができます。モデルのホストが成功すると、モデルは推論の準備が整うことになります。

SageMaker エンドポイントをテストする

SageMaker エンドポイントにモデルを無事にホストしたのち、空白の画像を使って動作検証をすることができます。テストコードは以下のようになります。

import numpy as np

ENDPOINT_NAME = 'yolov5l-demo'

modelHeight, modelWidth = 640, 640
blank_image = np.zeros((modelHeight, modelWidth, 3), np.uint8)
data = np.array(blank_image.astype(np.float32)/255.)
payload = json.dumps([data.tolist()])
response = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME,
    ContentType='application/json',
    Body=payload
)

result = json.loads(response['Body'].read().decode())
print('Results: ', result)

レイヤーによる Lambda 、ならびに起動トリガーの設定を行う

OpenCV を利用して画像データを送り、その推論結果を得ることでモデルが動くことを示します。Lambda には事前ビルドされた OpenCV のような外部ライブラリを搭載していないため、Lambda コードを呼び出す前にそれらをビルドしておく必要があります。さらに、Lambda が呼び出されるたびに OpenCV のような外部ライブラリをビルドすることは避けたいです。これらの目的のためLambda は Lambda レイヤーを作成する機能を提供します。レイヤーの中で動くものを定義することができ、Lambda コードが呼ばれるたびに使用することができます。ここでは OpenCV のための Lambda レイヤーの作成方法も示します。この記事では、Amazon Elastic Compute Cloud (Amazon EC2) をレイヤー作成に使います。

レイヤーが準備できたら、レイヤーで使用する Lambda コードを app.py スクリプトとしてを作成し、推論を実行して結果を取得します。このワークフローを図示すると以下のようになります。

Dockerを利用してOpenCVのLambda レイヤーを作成する

Python3.7 の Docker image を使用してDockerfileを作成すると、以下のようになります。

FROM amazonlinux

RUN yum update -y
RUN yum install gcc openssl-devel bzip2-devel libffi-devel wget tar gzip zip make -y

# Python 3.7のインストール
WORKDIR /
RUN wget https://www.python.org/ftp/python/3.7.12/Python-3.7.12.tgz
RUN tar -xzvf Python-3.7.12.tgz
WORKDIR /Python-3.7.12
RUN ./configure --enable-optimizations
RUN make altinstall

# Python パッケージのインストール
RUN mkdir /packages
RUN echo "opencv-python" >> /packages/requirements.txt
RUN mkdir -p /packages/opencv-python-3.7/python/lib/python3.7/site-packages
RUN pip3.7 install -r /packages/requirements.txt -t /packages/opencv-python-3.7/python/lib/python3.7/site-packages

# Lambda レイヤーデプロイのために zip ファイルの作成
WORKDIR /packages/opencv-python-3.7/
RUN zip -r9 /packages/cv2-python37.zip .
WORKDIR /packages/
RUN rm -rf /packages/opencv-python-3.7/

Dockerをビルドして処理を走らせ、layers ディレクトリ配下にZIPファイルを設置します。

$ docker build --tag aws-lambda-layers:latest <PATH/TO/Dockerfile>
$ docker run --rm -it -v $(pwd):/layers aws-lambda-layers cp /packages/cv2-python37.zip /layers

この処理の後、OpenCV レイヤーのアーティファクトをAmazon S3にアップロードし、Lambda レイヤーを作成します。

$ aws s3 cp layers/cv2-python37.zip s3://<BUCKET>/<PATH/TO/STORE/ARTIFACTS>
$ aws lambda publish-layer-version --layer-name cv2 --description "Open CV" --content S3Bucket=<BUCKET>,S3Key=<PATH/TO/STORE/ARTIFACTS>/cv2-python37.zip --compatible-runtimes python3.7

以上のコマンドを実行し処理が成功すると、Lambdaコンソール内にOpenCVのレイヤーが存在することを確認できます。

Lambda 関数を作成する

OpenCV を使用するためのスクリプトであるapp.pyを実行する Lambda 関数を作成します。
使用する際には次のコード BUCKET_NAME と IMAGE_LOCATION の値を、S3の画像にアクセスする場所に変更してください。

import os, logging, json, time, urllib.parse
import boto3, botocore
import numpy as np, cv2

logger = logging.getLogger()
logger.setLevel(logging.INFO)
client = boto3.client('lambda')

# S3 バケットの詳細
s3 = boto3.resource('s3')
BUCKET_NAME = "<NAME OF S3 BUCKET FOR INPUT IMAGE>"
IMAGE_LOCATION = "<S3 PATH TO IMAGE>/image.png"

# 推論エンドポイントの詳細
ENDPOINT_NAME = 'yolov5l-demo'
config = botocore.config.Config(read_timeout=80)
runtime = boto3.client('runtime.sagemaker', config=config)
modelHeight, modelWidth = 640, 640

# Lambdaの実行
def lambda_handler(event, context):
    key = urllib.parse.unquote_plus(event['Records'][0]['s3']['object']['key'], encoding='utf-8')

    # 入力 - 画像ファイルを Lambda の /tmp/ にダウンロード
    input_imagename = key.split('/')[-1]
    logger.info(f'Input Imagename: {input_imagename}')
    s3.Bucket(BUCKET_NAME).download_file(IMAGE_LOCATION + '/' + input_imagename, '/tmp/' + input_imagename)

    # 推論 - SageMaker エンドポイントの呼び出し
    logger.info(f'Starting Inference ... ')
    orig_image = cv2.imread('/tmp/' + input_imagename)
    if orig_image is not None:
        start_time_iter = time.time()
        # 入力画像の前処理
        image = cv2.resize(orig_image.copy(), (modelWidth, modelHeight), interpolation = cv2.INTER_AREA)
        data = np.array(image.astype(np.float32)/255.)
        payload = json.dumps([data.tolist()])
        # 推論の実行
        response = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME, ContentType='application/json', Body=payload)
        # 結果の取得
        result = json.loads(response['Body'].read().decode())
        end_time_iter = time.time()
        # 推論にかかった総時間の取得
        inference_time = round((end_time_iter - start_time_iter)*100)/100
    logger.info(f'Inference Completed ... ')

    # 出力 - 出力を下流のサービスで取り扱えるようにする
    return {
        "statusCode": 200,
        "body": json.dumps({
                "message": "Inference Time:// " + str(inference_time) + " seconds.",
                "results": result
                }),
            }

Lambda関数を以下のコードでデプロイします。

$ zip app.zip app.py
$ aws s3 cp app.zip s3://<BUCKET>/<PATH/TO/STORE/FUNCTION>
$ aws lambda create-function --function-name yolov5-lambda --handler app.lambda_handler --region us-east-1 --runtime python3.7 --environment "Variables={BUCKET_NAME=$BUCKET_NAME,S3_KEY=$S3_KEY}" --code S3Bucket=<BUCKET>,S3Key="<PATH/TO/STORE/FUNCTION/app.zip>"

Lambda関数にOpenCVレイヤーを取り付ける

Lambda 関数とが準備できたら、レイヤーと関数を以下のように接続します。

$ aws lambda update-function-configuration --function-name yolov5-lambda --layers cv2

レイヤーの設定はLambdaコンソールから確認することができます。

Amazon S3に画像がアップロードされたときにLambdaを起動させる

Amazon S3への画像アップロードをトリガーとしてLambda関数を起動します。この手順については、チュートリアル: Amazon S3 トリガーを使用して Lambda 関数を呼び出すを参照ください。

これらの手順を踏むことで、Lambda のコンソールで以下のように関数の詳細が表示されます。

推論を行う

LambdaとSageMaker エンドポイントのセットアップを行ったら、Lambda 関数を呼び出して出力をテストすることができます。Amazon S3への画像アップロードを、Lambda 関数の呼び出し、すなわち推論エンドポイントの呼び出しのトリガーとして使用します。例として、前のセクションで設定した Amazon S3 の場所 <S3 PATH TO IMAGE>/test_image.png に以下の画像をアップロードします。

画像のアップロードが行われると、Lambda 関数がトリガーされ、画像のダウンロードと読み込みを行い、推論を行うために SageMaker エンドポイントへ送られます。SageMakerエンドポイントからの出力は、その関数からJSONフォーマットで返ってくるので、様々な方法で利用することができます。以下は画像にオーバーレイされた出力例になります。

クリーンアップ

SageMaker ノートブックは、インスタンスタイプによっては、かなりの計算資源の使用とコストが必要になり得ます。不要なコストを排除するために、使用していないときには notebook インスタンスの停止を推奨します。また、Lambda 関数とは呼び出された場合にのみ課金される仕様です。そのため、これらのサービスのクリーンアップは不要です。しかしながら、SageMaker エンドポイントは ‘InService’ 状態で課金が発生するため追加コストを避けるために削除することを推奨します。

結論

この記事では、事前学習された YOLOv5 モデルを SageMaker エンドポイントでホストし、Lambda を利用して推論を呼び出し出力を処理する方法を紹介しました。詳細なコードなどは GitHub でも公開しています。

SageMaker エンドポイントの詳細については、「エンドポイントを作成してモデルをデプロイする」と「Amazon Sagemaker 推論モデルを構築、テストし、AWS Lambda にデプロイする」を参照してください。YOLOv5 モードのデプロイプロセスを自動化する方法が紹介されています。

Amazon Web Services ブログ