AWS Lambda がコンテナイメージをサポートしたので YOLOv5 を hosting してみた

どうも @kazuneet です。
初投稿です。こちらにも同様の内容を記載しております。

ことのあらまし

2020年末に AWS Lambda でコンテナイメージがサポートされました。
AWS Lambda の新機能 – コンテナイメージのサポート

この記事を見た AWS Lambda 至上主義であるあなたの上司がある日突然、

「今まで ML の推論環境として EC2/ECS/EKS/SageMaker を使って立てっぱなしにするのを黙認していたが、もう許さないよ？」
と言ってくることがあるかもしれません。

あるいはあなたの配偶者が、「つけっぱなしのリソースはもったいないと思わないんですか？」と言いながら「バッチ推論は起動時間がいつも遅いから嫌いっていっているでしょう？」とも言っている方もいるかもしれません。

そんな日に備えてコンテナイメージをサポートした AWS Lambda が ML の Just In Time の推論環境として使えるのかを今のうちに知っておくのは悪いことではありません。

やってみた

再掲ですが、基本的な使い方はこちらの記事が詳しいですし、そのとおりにやっていただければコンテナイメージで Node.js もしくは Python を動かす環境を作成できます。

AWS Lambda の新機能 – コンテナイメージのサポート

ですので、こちらの Python の例を改変して YOLOv5 を hosting してみようと思います。
（参考にしたものの、結果的にあまり原型はとどめておりません）

今回作成したコードは GitHub にアップロードしておりますので、みなさまぜひ clone して使っていただければ幸いです。

また、私が動かしている client 環境は SageMaker Notebook インスタンス(c5.xlarge)で、ロールに AdministratorAccess をｹﾞﾌﾝｹﾞﾌﾝ 適切なポリシーをアタッチしておりますので、適宜ご準備をお願いします。

準備

まずは YOLOv5 を AWS lambda ではなくお手元の環境で動かしてみましょう。といっても lambda を強く意識し、 yolov5/app.py に handler(event, context) というメソッドを用意して作ってみます。作成した handler メソッドを if __name__ == '__main__':から呼び出す形です。

AWS Lambda を WebAPI として使う方が多いかと思うので、event 引数では、imgキーで base64 encoding した画像データを受け取ることにします。

実際の推論を行うコードは、yolov5 の detect.pyから引数で受けている情報などを全てハードコーディングして切った貼ったしてデグレードしたわかりやすく固定値で埋め込み、推論部分を抜き出したものです。またオリジナルの detect.py は推論時にモデルがないと自動でダウンロードしてくるようになっているのですが、私が作成した GitHub は予めモデル(yolov5s.pt)も内包しております。

早速動かしてみましょう。

git clone https://github.com/kazuhitogo/yolov5-lambda
cd ./yolov5-lambda/yolov5
pip install -r requirements.txt
python app.py

すると current directory に./out.pngというのが出来上がります。これは、./data/images/bus.jpg を object detection したものです。

さて、handler が正常に動いていることを確認できたので、次にローカル環境でコンテナイメージをビルドし、ローカルの AWS Lambda 関数として動くかを確認してみます。

# git リポジトリのルートディレクトリへ
cd ../ 
# Build
docker build -t lambda-container-yolov5 .

この記事に則って（といっても大幅に改変していますが） Dockerfile を記載しておりますが、alpine に python-opencv が私の技術力じゃどうやっても入れられなかったのでライブラリの入れやすさの都合からイメージサイズの大きさには目を瞑って、 alpine ではなく slim-buster を利用しています。また YOLOv5 は Python 3.8 以上で動作するので、3.8 のイメージを利用しています。

さて、ローカルでコンテナを動かします。

docker run -p 9000:8080 lambda-container-yolov5:latest

さて、実際に post して動くかどうかを確認してみます。
同一インスタンスの別ターミナルから下記コマンドを実行してください。

# git リポジトリのルートディレクトリから
pip install -r requirements.txt
python test_code.py

すると、./out.png がgit リポジトリのルートディレクトリに出力されていますが、今度はジダン監督(?) などがdetection できていることがわかるかと思います。

ここで、気をつけないといけないのはローカルコンテナで動かす場合の endpoint は固定で、
http://localhost:9000/2015-03-31/functions/function/invocations であることに気をつけてください。今回は test_code.py にハードコーディングしてあります。

これで、ローカルでの確認は終わりですので、AWS Lambda にデプロイしてみましょう。

デプロイしてみる

build_and_test.ipunbの通りに実行していただければできますが、特記事項をいくつか。

まず、コンテナは他の処理と同様 ECR にプッシュすればOKです。

また、 lambda の実行ロールが必要なので、ロール作成後、適宜必要なポリシーをアタッチしておきます。

最後に、boto3 で lambda を create_function する際に、コンテナイメージを利用しない場合は zip ファイルの s3 URI を指定したりしましたが、コンテナイメージを利用する場合は、
ECRのURIにイメージのダイジェストを連結して指定する必要があります。逆に言えばそれだけで作成できます。また、コンテナイメージサイズが 3GB 程度あるので、初回呼び出し時にコンテナが動き出すまで時間がかかります。ですので、タイムアウトを 15 分最大の時間をとってあります。メモリサイズは適当に 1GB もあれば足りると思いますので、 1GB に指定しています。

create_function(
    FunctionName=function_name,
    Role=role_arn,
    Code={
        'ImageUri': f'{uri}@{image_digest}' # ECR の URI + ダイジェスト
    },
    Description='input-> b64img, output -> b64img, yolov5 detect',
    Timeout=60*15, # タイムアウト 15 分
    MemorySize=1024, # メモリ 1GB
    Publish=True,
    PackageType='Image',

さて、このセルでようやく推論しています。

使い方はいつもどおりですね。

%%time
res = lambda_client.invoke(
    FunctionName=function_name,
    Payload=json.dumps(data)
)

結果を画像に起こすとしっかりとジダン監督が見えるかと思います。いろんな画像に変えてやってみてください。

処理時間について

おそらく Lambda で ML の推論を利用する方は、処理時間が気になるのではないでしょうか？

SageMaker Hosting を利用した場合はレスポンスが良いですが、デプロイに数分かかってしましますし、ホスティングをしている間は課金されます。SageMaker のバッチ推論はやはりコンテナの起動時間がかかるので、わずかな推論の場合に数分のオーバーヘッドがかかるのはなんだかなぁ、と思うかと思います。

EC2 や ECS, EKS も似たような辛さがあるかと思います。

その点 AWS Lambda については、呼びたい時に短時間で呼び出せる はず で、動いている時間のみの課金のため、そういったワークロードでは威力を発揮するはずです。

あわせて、デプロイまでにどれくらい時間がかかるのかも見ておきたいところです。

というところで、東京リージョンの SageMaker Notebook インスタンス ml.c5.xlarge を利用し、この YOLOv5のアプリをデプロイした場合の 1 実験として、こうでした、というのをご参考まで。(jupyter のマジックコート %%time を利用。ネットワーク環境やコンピューティングリソースによって大きく変わるため、あくまでご参考）

イメージのビルドやプッシュ、関数の作成、初回実行については、alpine などの軽量イメージを使うなどで時間の短縮余地があるかと思います。

2回目以降の invoke については 2秒前半です。msec 単位の処理を求められる場合には厳しいですが、たまにデータが来て、10秒後に推論結果を出したい、といった場合にはかなり使える速度ではないでしょうか。また、おそらくメモリを増やしたらさらなる高速化が見込めます。

最後に

Lambda コンテナを使った ML 推論を試してみましたが、スケールし、かつ、秒単位での推論が求められるシーンでは使えそうなことがわかりました。また Step Functions などの ML パイプラインに組み込むことで、バッチ推論ほどじゃないモデルの簡単な精度検証にも使えるかもしれません。

これで、インスタンス立ちっぱなしやバッチ推論のオーバーヘッドが嫌いな上司や配偶者にも嫌われずにすむかもしれません。

Have a nice AWS Lambda Life !