Tensorflowを走らせるサーバーレスなWeb APIを建てる

:::note
この記事は元々**2024年8月23日**に個人ブログで公開された内容をQiitaに移行したものです。
:::

以前、Tensorflowの学習済みモデルを走らせるWeb APIを建てたいことがありました。

当時はシンプルにAWS EC2インスタンス上にPythonでWebサーバーを建てて、そこでTensorflowを実行していたのですが、これ結構無駄なお金がかかってしまっていたんです。  
というのも、Tensorflowを実行するにはそこそこのインスタンスのスペックが必要（nano, microではスペック不足で動かず、最低でもmini以上が必要）であり、さらに私が建てたかったWeb APIは高頻度で呼び出されることは想定されていませんでした。つまり、「安くないタイプのインスタンスを、そんなに呼ばれないのに24時間起動し続ける」というもったいないことをしていたんですね。

最近このWeb APIを、AWS Lambdaで建て直しましたので、この記事ではその方法を共有します。  
AWS LambdaでサーバーレスなAPIを建てれば、呼ばれた時だけ課金されるので、「呼ばれる頻度が高くない」という前提であればかなりコストを抑えることができます。

## テンプレート用意

今回は私が以前作った「AWS Lambdaで楽にWeb APIを建てるためのテンプレ」を使っていきます。

https://github.com/ShimeiYago/lambda-api-template

このテンプレを使って新しいリポジトリを作成し、お手元に`git clone`しておいてください。

![git-template.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3981597/a030ec89-f21c-4b7b-9b54-d179f53f749e.jpeg)

クローンしたら、リポジトリのルートに移動しておいてください。

※このテンプレを使ったAPIを建てる手順はすべて以下の記事で解説していますので、まずはこの記事を参照しながらサクッとWeb APIを建ててみてください。

[link]

上記の手順で無事APIを建てることができたなら、あとはこのコードを一部変更して、Tensorflowを走らせるAPIに改造していきます。

## Tensorflowの公式イメージをベースにする

テンプレの`Dockerfile`は`ubuntu:20.04`をベースイメージにしていますが、今回はTensorflowを使う目的で、Tensorflowの公式イメージをベースにしようと思います。`Dockerfile`の１行目を次のように書き換えます。

```dockerfile
FROM tensorflow/tensorflow:2.16.2
```

Tensorflowのバージョンは好きなものを選択してください。

## ライブラリを追加する

`requirements.txt`を編集し、必要なライブラリを追加します。今回は例としてPillowをインストールしますが、ここではアプリケーションに必要なライブラリを自由に記述してください。

```text
awslambdaric # This is mandatory
pillow==10.4.0
```

Tensorflowのベースイメージを使っているので、Tensorflowはすでにインストール済みです。ここではTensorflowを記載する必要はありません。

## 学習済みモデルを用意する

学習済みモデルを別途用意して、Build時にDockerイメージの中に組み込みます。

今回は例として、VGG16の学習済みモデルを使用します。VGG16はImageNetデータセットで学習された画像分類モデルで、入力画像を1000クラスに分類することができます。

VGG16の学習済みモデルは一般公開されているので、ダウンロードします。

```bash
mkdir models
```

```bash
curl -Lo models/vgg16_weights_tf_dim_ordering_tf_kernels.h5 https://storage.googleapis.com/tensorflow/keras-applications/vgg16/vgg16_weights_tf_dim_ordering_tf_kernels.h5
```

加えて、結果表示のために画像クラスの名前リストもダウンロードしておきます。（例えば１番目のクラスは「ドクターフィッシュ」、２番目は「ゴールドフィッシュ」といった情報ですね。）

```bash
curl -Lo models/imagenet_class_index.json https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json
```

最後に、これらのファイルをDockerイメージに組み込むために、`Dockerfile`を以下のように更新します。

```dockerfile
FROM tensorflow/tensorflow:2.16.2

RUN apt-get update && apt-get install -y python3-pip

COPY ./requirements.txt /requirements.txt
RUN pip install -r /requirements.txt

ARG APP_DIR="/app"

RUN mkdir ${APP_DIR}

COPY app.py ${APP_DIR}

COPY ./entry_script.sh /entry_script.sh
ADD ./aws-lambda-rie /usr/bin/aws-lambda-rie

# ここから
ARG MODEL_DIR="${APP_DIR}/models"
RUN mkdir -p ${MODEL_DIR}
COPY models/vgg16_weights_tf_dim_ordering_tf_kernels.h5 ${MODEL_DIR}
COPY models/imagenet_class_index.json ${MODEL_DIR}
ENV TF_CPP_MIN_LOG_LEVEL="2"
# ここまで追加

WORKDIR ${APP_DIR}

ENTRYPOINT [ "/entry_script.sh" ]

CMD [ "app.handler" ]
```

## app.pyを書き換える

実際にどのような処理を行い、どのようなレスポンスを返すのかは`app.py`に記述します。今回はWeb APIのリクエストとして画像データを受け取り、レスポンスとしてVGG16による分類結果を返すような処理を記述していきます。

`app.py`を以下のように更新します。

```python
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras import preprocessing
import numpy as np
import base64
import PIL
import io
import json

def handler(event, context):
    # 学習済みモデルをロード
    model = VGG16(weights='models/vgg16_weights_tf_dim_ordering_tf_kernels.h5')

    # リクエストBodyから入力画像(base64形式)を取得
    input_image_base64 = event["inputImageBase64"]
    
    # 入力画像をPIL形式に変換
    img = base64_to_pil(input_image_base64)

    # サイズをVGG16のデフォルトである224x224にリサイズ
    img = img.resize((224, 224))

    # 読み込んだPIL形式の画像をarrayに変換
    x = preprocessing.image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    
    # 予測
    preds = model(preprocess_input(x))
    
    # 最も高い確率のクラス名を取得
    max_index = np.argmax(preds)
    with open("models/imagenet_class_index.json") as f:
        class_index = json.load(f)
        class_name = class_index[str(max_index)][1]
    
    # レスポンス
    res = {
        "class": class_name
    }
    return res

# base64画像文字列をデコードしてPIL画像に変換
def base64_to_pil(img_base64):
    img_data = base64.b64decode(img_base64)
    img = PIL.Image.open(io.BytesIO(img_data))
    return img
```

## DockerイメージをBuildする

改めてDockerイメージをbuildします。

```bash
docker build -t lambda-api .
```

## ローカルでテストする

ローカルでDockerコンテナを起動してテストします。

```bash
docker run --rm -p 8000:8080 lambda-api:latest
```

Web APIが[localhost:8000/2015-03-31/functions/function/invocations](http://localhost:8000/2015-03-31/functions/function/invocations) でアクティブになります。

テスト用に、何か適当な画像を送信してみましょう。

パブリックドメインのワンちゃんの画像を使います。

![publicdomainq-0081090uok.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3981597/f0660d4f-a804-4ce2-9685-d7af345bce2c.jpeg)


この画像を、[こちらのサービス](https://base64.guru/converter/encode/image)などを利用してBase64形式の文字列に変換して、以下の形式のリクエストでWeb APIを叩きます。

```json
{
  "inputImageBase64": "/9j/..."
}
```

するとVGG16によってクラスが予測され、結果が以下のように返ってきます。

```json
{
  "class": "beagle"
}
```

「beagle」とは「ビーグル犬」のことなので、正しく分類できていることがわかります。

ここまで来ればあとはLambdaにデプロイするだけなので、[テンプレの手順](https://tech.dbstory.jp/docker-python-lambda-api-template/)に従ってサクッとデプロイしちゃいましょう。

## 注意点：Lambdaのタイムアウトとメモリ

AWS Lambdaはデフォルトだとタイムアウトが3秒でメモリが128MBですが、これだとTensorflowを動かすには不足です。

![lambda-setting.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3981597/91cb884a-2f68-4da0-b15b-de5b350efaf4.png)

タイムアウトに関してですが、Tensorflowは実行時間がかかるものです。モデルの大きさにもよりますが、**30秒以上**はあったほうが良いと思います。

メモリに関しても、Tensorflowは結構メモリを食います。少なくとも1000MBはあったほうが良いと思いますが、これもモデル次第でもっとあったほうが良さげです。多めに設定してみて、実際のメモリ利用を見てみて、ちょうど良い値まで減らすと良いでしょう。

しかもAWS Lambdaでは、CPU数はメモリの大きさに依存します。大きいメモリを設定すればCPU数も増えるのです。  
TensorflowをCPUで動かすためには多めのCPU数が必要です。もちろんメモリを多く設定すると料金も増えますが、CPU数が増えることが実行時間の削減に繋がるので、結果的に料金の節約になります。

Lambdaのメモリ数とCPU数の関係は以下の記事にまとめられています。私の場合はTensorflowを実装する時は3CPUはほしいので、3CPUが出る**3072MB**を設定しています。

https://qiita.com/komikoni/items/101ccd3cb98fdc4f4389

![lambda-setting-tf-1024x177.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3981597/ce0d893a-2f38-4c43-905a-749b0897850a.jpeg)