こんにちは、R&Dの宮崎です。普段はTensorFlowを使った画像の認識モデルの開発を行なっています。無事に精度の高いモデルができると、次は実際にサービスとして運用するための基盤を準備しなければいけません。そこで今回は、AWS LambdaでTensorFlow 2.0を動かせるか検証してみましたので、紹介したいと思います。

TensorFlowで推論するための基盤

TensorFlowの運用基盤としてはいくつかの選択肢が考えられます。 AWSではTensorFlow ServingのコンテナをECSやEKSで動かしたり、SageMakerを使ったりするなどです。これらに対し、AWS Lambdaはサーバのプロビジョニングが不要で運用の手間がかからず、使った時間しか課金が発生しないため、推論リクエストが散発的に発生する場合はコストを抑えられるといったメリットがあります。一方で、AWS Lambdaはコールドスタートからの起動時間がかかったり、使えるコンピュータリソースが少ないというデメリットもあります。特に、TensorFlowを扱う上ではその大きなパッケージやモデルを、いかに限られたディスクサイズに収めるかが鍵となります。本記事ではサイズの削減方法を中心に検証していきます。

モデルの作成

まず、画像分類モデルを作成します。今回はこちらのサンプルをベースにResNet50を使ってImageNetで学習した学習済みモデルを使用します。

from tensorflow.keras.applications.resnet50 import ResNet50

model = ResNet50(weights='imagenet')
model.save('/path/to/model_resnet50_imagenet.h5', 
           include_optimizer=False)

103MBの画像分類モデルが作成されました。

AWS Lambdaへのデプロイ

デプロイに使用するソフトウェア

AWS LambdaにデプロイするためのツールとしてServerlessを使用します。 AWS Lambdaは基本的なPythonパッケージしかインストールされていないため、TensorFlowなどは自分で追加インストールする必要があります。そこで、追加パッケージのインストールに便利なserverless-python-requirementsプラグインもあわせて使用します。

ソフトウェア	バージョン
serverless	1.52.0
serverless-python-requirements	5.0.0

いかにディスクサイズに収めるか

AWS Lambdaで使えるディスクサイズはPythonコードを動かすためのワーキングディレクトリが250MB, 一時利用可能な/tmpディレクトリが512MBです。一方でTensorFlow 2.0は依存パッケージも含めると498MBのため、250MBをオーバーしてしまいます。そこで、serverless-python-requirementsに含まれている、不要ファイルを削除し、ZIP圧縮して/tmpディレクトリに展開する機能を使用します。またTensorFlow依存パッケージのうち、Numpyなどの一部のパッケージはLambda Layerを使用し、ワーキングディレクトリに配置します。これにより、/tmpディレクトリを使い切らず、推論モデルを配置するための容量確保を図ります。今回の推論モデルは103MBですが、複雑なネットワークを使用したり、前処理・後処理も含めたりするとモデルサイズも大きくなります。この配置であれば250MB弱までのモデルなら載せることができます。

f:id:unifa_tech:20190912113350p:plain — ディレクトリ毎のファイル配置

s3バケットの準備とモデルファイルのアップロード

事前にs3バケットを作成し、推論モデルファイルをアップロードします。今回はimage-classification-lambdaという名前でバケットを作成しました。

Serverless定義のディレクトリ構成

Lambda関数とLambda Layerの2つのディレクトリを、以下のようにローカル環境に作成します。

├── lambda-function
│   ├── handler.py
│   ├── requirements.txt
│   └── serverless.yml
└── lambda-layer
    ├── requirements.txt
    └── serverless.yml

Lambda関数

まずはlambda-functionディレクトリ内にLambda関数を定義します。

handler.py

Lambda関数のハンドラを記述するためのpythonコードです。

# (1) TensorFlowパッケージの展開
try:
    import unzip_requirements
except ImportError:
    pass

import io
import os
from boto3.session import Session
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 \
    import preprocess_input, decode_predictions
import numpy as np

MODEL_BUCKET ='image-classification-lambda'
MODEL_NAME = 'model_resnet50_imagenet.h5'
MODEL_PATH = os.path.join('/tmp', MODEL_NAME)

# (2) モデルのダウンロード
session = Session()
s3_client = session.client('s3')
s3_client.download_file(MODEL_BUCKET, MODEL_NAME, MODEL_PATH)

model = None


def classify(event, context):
    global model
    if model is None:
        model = load_model(MODEL_PATH, compile=False)

    # (3) 画像のダウンロード
    s3_object = s3_client.get_object(
        Bucket=event['bucket'], Key=event['filename'])
    image_data = io.BytesIO(s3_object['Body'].read())
    img = image.load_img(image_data, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

    # (4) 推論
    preds = model.predict(x)
    results = decode_predictions(preds, top=3)[0]

    return [{'class': r[1], 'score': float(r[2])} for r in results]

まず(1)でserverless-python-requirementsがZIP圧縮してくれたTensorFlowパッケージを展開します。(2)で事前にS3に設定していた推論モデルファイルをダウンロードします。(3)で推論リクエストのあった画像ファイルをダウンロードし、(4)で推論します。

serverless.yml

SeverlessでAWS Lambdaにデプロイするための定義ファイルです。

service: image-classification

plugins:
  - serverless-python-requirements

custom:
  pythonRequirements:
    dockerizePip: true
    zip: true
    slim: true
    slimPatterns:
      - "**/debug"
      - "**/grpc"
      - "**/h5py"
      - "**/markdown"
      - "**/numpy"
      - "**/pkg_resources"
      - "**/setuptools"
      - "**/tensorboard/plugins"
      - "**/tensorboard/webfiles.zip"
      - "**/tensorflow_core/contrib"
      - "**/tensorflow_core/examples"
      - "**/tensorflow_core/include"
      - "**/tensorflow_estimator"
      - "**/werkzeug"
      - "**/wheel"
  requirementsService: image-classification-layer
  requirementsExport: ImageClassificationLayer
  requirementsLayer: ${cf:${self:custom.requirementsService}-${self:provider.stage}.${self:custom.requirementsExport}}

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

functions:
  classify:
    handler: handler.classify
    memorySize: 2048
    timeout: 60
    layers:
      - ${self:custom.requirementsLayer}

ポイントはcustom.pythonRequirements.zipをtrueにし、ZIP圧縮を有効化します。そして、custom.pythonRequirements.slimPatternsにTensorFlow依存パッケージのうち使用しないファイルやLambda Layerに載せるファイルなど削除するパスをひたすら書くことで498MBから264MBまで容量を削減します。また、functions.classify.layersにこの後に定義するLambda Layerを指定します。

requirements.txt

tensorflow==2.0.0-rc0

requirements.txtにはserverless-python-requirementsにZIP圧縮してもらうためのパッケージを記載します。

Lambda Layer

次にlambda-layerディレクトリ内にLambda Layerを定義します。

serverless.yml

service: image-classification-layer
plugins:
  - serverless-python-requirements

custom:
  pythonRequirements:
    dockerizePip: true
    layer:
      compatibleRuntimes:
        - python3.6
    slim: true
    strip: false
    slimPatterns:
      - "**/tests"

provider:
  name: aws
  runtime: python3.6
  stage: dev
  region: ap-northeast-1

resources:
  Outputs:
    ImageClassificationLayer:
      Value:
        Ref: PythonRequirementsLambdaLayer

serverless-python-requirementsにおいて、custom.pythonRequirements.slimをtrueとすると、デフォルトではsoファイルからシンボル情報を取り除いてサイズを削減してくれます。しかし、h5pyやPillowなど一部のパッケージではエラーを引き起こしてしまいます。そこで、それらのパッケージはシンボル情報を削除しないようcustom.pythonRequirements.stripをfalseにしてLambda Layerに格納します。

また、custom.pythonRequirements.slimPatternsにはNumpyのテストコードを削除するためのパスを記載します。

requirements.txt

h5py==2.10.0
numpy==1.17.2
Pillow==6.1.0

ZIP圧縮しない(Lambda関数に含めない)容量の大きいNumpyパッケージとシンボル情報を削除したくないh5pyとPillowパッケージを書きます。

デプロイ

設定したLambda LayerとLambda関数をAWS Lambdaにデプロイします。

$ cd lambda-layer
$ serverless deploy
$ cd ../lambda-function
$ serverless deploy

これで準備が整いました。

推論の実施

以下の象の画像を使って推論してみましょう。リクエスト前に先ほど作成したimage-classification-lambdaバケットにelephant.jpgとして格納しておきます。

f:id:unifa_tech:20190912133323j:plain — 推論リクエスト画像 (出典: Open Images Dataset V5)

それではデプロイしたLambda関数に推論リクエストを投げます。

$ serverless invoke \
    --function classify \
    --data '{"bucket": "image-classification-lambda", "filename": "elephant.jpg"}'
[
    {
        "class": "Indian_elephant",
        "score": 0.830479621887207
    },
    {
        "class": "tusker",
        "score": 0.15385641157627106
    },
    {
        "class": "African_elephant",
        "score": 0.015656592324376106
    }
]

無事に分類してくれました。どうやらこの象はインド象のようです。

おわりに

今回はAWS LambdaでTensorFlow 2.0を動かしてみました。深層学習の技術はまだまだ発展途上のため、ライブラリやモデルの構築方法、運用基盤など様々な選択肢が存在します。その中から要件にあった適切な技術を選んでいきたいと思います。

ユニファ開発者ブログ

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

AWS LambdaでTensorFlow 2.0を使った画像分類

TensorFlowで推論するための基盤

モデルの作成

AWS Lambdaへのデプロイ

デプロイに使用するソフトウェア

いかにディスクサイズに収めるか

s3バケットの準備とモデルファイルのアップロード

Serverless定義のディレクトリ構成

Lambda関数

handler.py

serverless.yml

requirements.txt

Lambda Layer

serverless.yml

requirements.txt

デプロイ

推論の実施

おわりに