AWS LambdaのコンテナデプロイでYOLOv5の推論エンドポイントを作成する

データアナリティクス事業本部 インテグレーション部 機械学習チームの貞松です。

本記事は「クラスメソッド 機械学習チーム アドベントカレンダー 2022」の9日目です。
昨日(8日目)は Shirota による 日本語自然言語処理オープンソースライブラリ「GiNZA」で構文解析をやってみた でした。

AWS Lambda(以降Lambda)は、様々な処理を関数単位でサーバーレスに実行可能なサービスです。
気軽に処理を実装、実行でき、複数の処理をイベント駆動で制御することによって、あらゆるシステムのバックエンドとして機能します。
非常に便利なので私もよく利用するのですが、最近最早「全部Lambdaで処理させれば良いんじゃない？」という風潮を感じています。

本記事では、思い切って機械学習の推論処理をLambdaに載せてしまおうという試みで、物体検出アルゴリズムであるYOLOv5の処理をLambdaのコンテナデプロイを利用して実装します。

処理の実装

AWS CDKプロジェクトの作成

まず、Lambda含めた諸々のリソースをデプロイするためのAWS CDKプロジェクトを作成します。

cdk init app --language=typescript

YOLOv5のソースを取得

YOLOv5のソースコードを下記のGitHubリポジトリから取得して、Lambdaと同じ階層のディレクトリに配置します。

Lambda関数を作成

YOLOv5のソースコードと学習済みモデルを使って画像中の物体検出を実行する処理を作成します。
Lambda関数は、ローカルのファイル入出力の処理に tmp 以外の領域を使用できないことに十分注意してください。
tmp 配下にサブディレクトリを作成して操作することもできません。

lambda/yolo_v5s_inference.py

import os
import sys
import json
import logging
import torch
import boto3
from pathlib import Path
module_path = os.path.abspath(os.path.join("./yolov5"))
if module_path not in sys.path:
    sys.path.append(module_path)
from models.experimental import attempt_load
from utils.dataloaders import LoadImages
from utils.general import check_img_size, non_max_suppression, scale_boxes
from utils.plots import save_one_box
from utils.torch_utils import select_device

# # YOLOの処理に必要なパラメータ
IMAGE_SIZE = 1200
CONF_THRESH = 0.4
IOU_THRESH = 0.45
# 学習済みモデルファイルの配置パス
MODEL_PATH = "./yolov5/yolov5s.pt"
# S3リソースの生成
S3 = boto3.resource("s3")
BUCKET = S3.Bucket("cm-sadamatsu-resource-20221209")

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

def inference():
    # Lambda上で実行するためCPU実行
    device = select_device('cpu')

    # load_model
    model = attempt_load(MODEL_PATH, device=device)
    model.info(False, IMAGE_SIZE)

    stride = int(model.stride.max())
    imgsz = check_img_size(IMAGE_SIZE, s=stride)
    img_path = "/tmp/image.jpg"
    dataset = LoadImages(img_path, img_size=imgsz, stride=stride)

    for path, img, im0s, vid_cap, s_count in dataset:
        img = torch.from_numpy(img).to(device).float()
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        pred = model(img, augment=True)[0]

        # Apply NMS
        pred = non_max_suppression(
            pred, CONF_THRESH, IOU_THRESH, classes=None, agnostic=False)

        detect_image_count = 0
        for i, det in enumerate(pred):  # detections per image
            p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
            p = Path(p)  # to Path
            imc = im0.copy() # for save_crop
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], im0.shape).round()
                # Write results
                for *xyxy, conf, cls in det:
                    save_image_path = f"/tmp/crop_image_{detect_image_count}.jpg"
                    save_one_box(xyxy, imc, file=Path(save_image_path, BGR=False))
                    # S3 にアップロードする
                    image_s3_key = f"yolov5/output/crop_image_{detect_image_count}.jpg"
                    BUCKET.upload_file(save_image_path, image_s3_key)
                    detect_image_count += 1

def handler(event, context):
    try:
        BUCKET.download_file("yolov5/input/image.jpg", "/tmp/image.jpg")
        inference()

        return {
            "statusCode": 200,
            "body": json.dumps("Success yolov5s inference execution.")
        }
    except Exception as ex:
        logger.error(ex)
        return {
            "statusCode": 500,
            "body": json.dumps("Failed yolov5s inference execution.")
        }

CDKのソースコードを作成

Lambda関数をコンテナデプロイする為に必要なリソース(実行権限やコンテナイメージの取得元になるECRなど)を一式デプロイするためのCDK用ソースコードを作成します。

lib/pytorch_hub_yolov5_on_lambda-stack.ts

#!/usr/bin/env node
import 'source-map-support/register';
import { App } from 'aws-cdk-lib';
import { PytorchHubYolov5OnLambdaStack } from '../lib/pytorch_hub_yolov5_on_lambda-stack';

const app = new App();
new PytorchHubYolov5OnLambdaStack(app, 'PytorchHubYolov5OnLambdaStack', {});

bin/pytorch_hub_yolov5_on_lambda.ts

import { Stack, StackProps, Duration, Tags } from "aws-cdk-lib";
import { Construct } from "constructs";
import { aws_iam as iam } from "aws-cdk-lib";
import { aws_lambda as lambda } from "aws-cdk-lib";
import { aws_ecr as ecr } from "aws-cdk-lib";

export class PytorchHubYolov5OnLambdaStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    // ECRリポジトリARNを指定して、Lambda関数用カスタムコンテナイメージのリポジトリを取得
    const lambdaEcrRepository = ecr.Repository.fromRepositoryArn(
      this,
      id,
      "arn:aws:ecr:ap-northeast-1:xxxxxxxxxxxx:repository/yolov5-lambda",
    );

    // コンテナイメージを利用するLambda関数で必要な実行ロールを設定
    const execLambdaRole = new iam.Role(this, "execRole", {
      roleName: `cm-sadamatsu-lambda-exec-role`,
      assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
      managedPolicies: [
        iam.ManagedPolicy.fromAwsManagedPolicyName(
          "service-role/AWSLambdaBasicExecutionRole"
        ),
        iam.ManagedPolicy.fromAwsManagedPolicyName(
          "AmazonEC2ContainerRegistryReadOnly"
        ),
      ],
    });

    // 開帳幅推定処理のLambda関数
    const fnYolov5sInference = new lambda.Function(this, "yolo_v5s_inference", {
      code: lambda.Code.fromEcrImage(lambdaEcrRepository, {
        cmd: ["yolo_v5s_inference.handler"],
        tag: "latest",
        entrypoint: ["/lambda-entrypoint.sh"],
      }),
      role: execLambdaRole,
      functionName: "yolo_v5s_inference",
      runtime: lambda.Runtime.FROM_IMAGE,
      handler: lambda.Handler.FROM_IMAGE,
      memorySize: 2048,
      timeout: Duration.seconds(600),
    });
    Tags.of(fnYolov5sInference).add("runtime", "python");

    // S3バケットにアクセスする為の権限(ポリシー)を作成してアタッチ
    const bucketAccessPolicy = new iam.PolicyStatement({
      resources: ["*"],
      actions: ["s3:PutObject", "s3:GetObject", "s3:ListBucket"],
    });
    fnYolov5sInference.addToRolePolicy(bucketAccessPolicy);
  }
}

コンテナイメージの作成

Lambdaで使用する為のコンテナイメージを作成します。

requirements.txtの作成

Dockerfile内でパッケージをインストール処理部分に使用するrequirements.txtを作成します。
コンテナデプロイされたLambda上でYOLOv5を実行するために必要なパッケージをひたすら詰め込んでいます。

opencv-python
opencv-contrib-python
pandas
scikit-learn
requests
IPython
Pillow
psutil
pyyaml
tqdm
matplotlib
seaborn

Dockerfileの作成

Dockerfileを作成します。
YOLOv5でサポートされているライブラリバージョンの考慮と、LambdaがGPU対応していない事を考慮して、torchおよびtorchvisionパッケージはrequirements.txtには含めず、Dockerfile内で別途インストールの処理を記述しています。 また、Lambdaのコンテナデプロイ用に作成する為に、パッケージのインストールパスやワークディレクトリパス、オブジェクトの実行権限等のお作法を守る必要があるので十分注意してください。

# Lambda関数用のカスタムコンテナイメージのDockerfile
FROM public.ecr.aws/lambda/python:3.8

RUN yum install -y mesa-libGL cmake

COPY lambda/requirements.txt /opt/

RUN pip install -U pip
RUN pip install -U setuptools
RUN pip install -U cython
RUN pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html -t /opt/python/lib/python3.8/site-packages
RUN pip install -r /opt/requirements.txt -t /opt/python/lib/python3.8/site-packages

#Function code
WORKDIR /var/task
COPY lambda/ .
RUN find /var/task -type f -exec chmod 644 {} +
RUN find /var/task -type d -exec chmod 755 {} +

コンテナイメージのビルドとリポジトリへのプッシュ＆cdk deployを実行

コンテナイメージをビルドしてリポジトリにプッシュ→cdk deployまで実行するシェルスクリプトを作成しました。
処理の最後で、Lambda関数の処理を修正・更新することを考慮したLambda関数の参照をリフレッシュして最新のコンテナイメージを参照させる処理を入れています。
コンテナイメージのプッシュ先リポジトリは予めECRで適当に作成して、リポジトリの名称をスクリプト実行時に渡します。

#!/bin/bash

ACCOUNTID=${1}
REGION=${2}
ECR_REPO_NAME=${3}

ECR_REPO_URI="${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/${ECR_REPO_NAME}"

echo "ECR_REPO_URI: ${ECR_REPO_URI}"

aws ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin ${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com
docker build -f ./docker/dockerfile -t ${ECR_REPO_NAME} .
docker tag ${ECR_REPO_NAME}:latest ${ECR_REPO_URI}:latest
docker push ${ECR_REPO_URI}:latest

cdk deploy -c ecrRepoArn=arn:aws:ecr:${REGION}:${ACCOUNTID}:repository/${ECR_REPO_NAME}
aws lambda update-function-code --function-name yolo_v5s_inference --image-uri ${ECR_REPO_URI}:latest

最終的なプロジェクト構成

最終的なプロジェクトのディレクトリ構成は以下のようになりました。

デプロイしたLambda動作確認

S3バケットと入力画像の準備

今回は以下のような構成でS3バケットを用意しておきます。

cm-sadamatsu-resource-20221209
  └　yolov5
    ├ input
    │  └　image.jpg
    └ output

入力画像のimage.jpgは、YOLOv5用のサンプル画像でおなじみ、zidane.jpgをリネームして配置しておきます。

Lambda関数の実行

Lambda関数を実行すると、検出した領域を切り出して　cm-sadamatsu-resource-20221209/output/yolov5/output 配下に出力します。

切り抜かれた画像は以下の通りです。

person(ジダン監督)

tie(ジダン監督のネクタイ)

person(アンチェロッティ監督)

アンチェロッティ監督のネクタイが検出できていませんが、これは物体検出処理時にconfidenceの閾値を0.4で設定した為です。
※今回使用した学習済みのyolov5sモデルを使用すると、アンチェロッティ監督のネクタイのconfidenceは0.26になる模様です。

まとめ

以上、Lambdaのコンテナデプロイを使用したYOLOv5の推論エンドポイント作成方法の解説でした。
今回はとにかく動かすところまでに留めているので、まともにやろうとすると少々実装コスト高めですね。
Lambdaの制約等々、格闘する部分が多く、突貫工事になっている部分も多い為、あくまでこんなこともできるよ、という一例として捉えていただけると幸いです。
もっとスマートな実装を考案しつつ、AWSサービス自体のアップデートにも期待したいところですが、差し当たりLambdaのGPU対応お待ちしております。