AWS LambdaでディープラーニングAPIを実装（AWS Lambda + Amazon EFS）

TensorFlow (Python)で学習したモデルをAWS LambdaのAPIに乗っけて簡単に呼び出したいと考えたのですが、実装時にいろいろとつまずいたので備忘録として記録したいと思います。

また、学習済みのモデルをReactのフロントエンドで動かす方法はこちらで紹介しています。

AWS LambdaでTensorFlowを実装する主な方法

TensorFlowのようなディープラーニング用ライブラリはファイルサイズが大きいため、そのままzipで固めてアップロードしようとしてもAWS Lambdaのアップロード制限に引っかかってしまいます。そのため、ライブラリをうまく配置する必要があります。

TensorFlowの配置方法	特徴
ライブラリ内の必要な要素だけを抽出してzipにし，AWS Lambdaの/tmp領域に展開(参考)	AWS Lambda, Amazon S3以外のサービスを利用する必要がない。/tmp領域に一度TensorFlow, 推論モデルを展開するため、モデルサイズによっては/tmpの500MBの制限を超えてしまう。
TensorFlowをEFSに配置	EFSの領域サイズはほぼ無制限なのでAWS Lambdaのメモリ制限さえ注意すればどんなサイズのモデルでも扱える。EFSにTensorFlowをインストールするためにEC2上でのインストール作業が必要。
TensorFlow実行環境を持つDockerコンテナでVPC Lambdaに配置(公式)	EFS同様領域サイズはほぼ気にする必要はない。VPC LambdaなのでVPCが必要。

今回は二つ目のEFSによるやり方を説明します。

Amazon EFSへTensorFlowをインストール

EFSを設定

こちらを参考にEFSを設定します。

EFSをEC2にマウント

EFSが設定できたら、consoleのEFSから作成したファイルシステムを選択し、右上の「アタッチ」をクリックします。するとマウントするためのコマンドが出てくるのでこれをコピーしておきます。

次にEC2にアクセスし、EFSをマウントするディレクトリを作成します。(rootで実行)

$ sudo su -
# mkdir ~/efs # マウント先のディレクトリ
# yum install -y amazon-efs-utils # マウントヘルパーをインストール
# mount -t efs -o tls fs-xxxxxxx:/ ~/efs # console上で記載されていたコマンドをコピペ（マウント先は作成したディレクトリ）

これで作成したEFSをEC2にマウントすることができました。

EC2からtensorflowをインストール

最後にマウントしたEFSにTensorFlowをインストールします。(同様にrootで実行)

# pip3 install tensorflow==2.4.0 -t ~/efs/{EFSで設定した名前}

ここで注意しておきたいのは、このコマンド上のpip3のPythonのバージョンとAWS Lambda上でのPythonのバージョンが一致しているかです。バージョンが一致していないとTensorFlowが起動できません。もしバージョンが一致していない場合はpyenv等で仮想環境を用意してインストールを行います。

Serverless FrameworkによるAWS Lambdaデプロイ

今回はEFSを介してTensorFlowを起動するところまでを行うため、コードは以下の通り最小の構成になっています。Serverless Frameworkの設定はGithubを参考にしてください。

ポイントはsys.path.appendでEFS上のTensorFlowをimportできるようにするところです。

import os
import sys
import json
import logging
sys.path.append(os.environ.get('EFS_MOUNT_PATH'))
import tensorflow as tf

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

def predict(event, context):
    return {
        'statusCode': 200,
        'body': json.dumps({
            'tf_version': tf.__version__
        })
    }