AWSLambdaでPyTorch【Lambda import編】

# はじめに
本記事はAWSのLambda上でPyTorchを動かして見ようという試みについてのまとめです。DeepLearningタグをつけていますが、学習については触れません。ゴールはLambda上で何かしらの推論を動かしてみるというところまでです。（EFS設定編は[こちら](https://qiita.com/myonoym/items/5d360ed11931474cb640)）

# ざっくりまとめ
1. EFS使ってみる
1. LambdaでPyTorch←今回ココ
1. slackからも呼ぶ

の三本立てです。

# Lambdaで動かすモデル
今回Lambdaで動かしてみるのは、[EML-NET](https://www.sciencedirect.com/science/article/abs/pii/S0262885620300196)です。[Githubのプロジェクトページ](https://github.com/SenJia/EML-NET-Saliency)を参照すると詳細が書いてありますが、[SaliencyMap](https://en.wikipedia.org/wiki/Saliency_map)を生成するモデルです。SaliencyMapとは人の視線の向く位置をヒートマップで表現したものになります。

[論文Figure1より引用](https://arxiv.org/pdf/1805.01047.pdf)
![スクリーンショット 2020-09-14 17.16.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/575891/2394a07f-b155-2c3c-ae7c-88e13d017688.png "Figure1")


# pipで色々入れる
[前回の記事](https://qiita.com/myonoym/items/5d360ed11931474cb640)で/mntにEFSをマウントしたのですが、EFSにLambdaからimportしたいライブラリをポンポコ入れていきましょう。

```sh
$ cd /mnt/lambda
$ sudo pip3 install -t . torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
$ sudo pip3 install --upgrade -t . opencv-python==4.4.0.42 scipy==1.5.2
```
これによって、

- torch==1.6.0+cpu
- torchvision==0.7.0+cpu
- numpy==1.19.2
- future==0.18.2
- pillow==7.2.0
- opencv==4.4.0.42

が`/mnt/lambda`以下に入りました。

# Lambdaでimport
## Lambdaを用意する
`関数の作成` -> `一から作成`をチェック -> 関数名を入力 -> ランタイムを`Python3.7`にする -> `関数の作成`をクリックしましょう

## VPC
[このあたり](https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/configuration-vpc.html)になにやら書いてありますが、LambdaからEFSに接続するためにはLambdaをVPC内に配置する必要があります。デフォルトのVPCで大丈夫でしょう。セキュリティグループに関しては、EFS内にライブラリを用意するために使ったEC2の設定があると思うのでそちらを指定しておきます。
保存をクリックすると、以下のようなエラーが出てしまいます。
![スクリーンショット 2020-09-15 12.04.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/575891/f78d7b97-273a-ef78-0828-e57e4e36a6b4.png)
権限が足りないようです。[こちら](https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/configuration-vpc.html#vpc-permissions)によると、

```
ec2:CreateNetworkInterface
ec2:DescribeNetworkInterfaces
ec2:DeleteNetworkInterface
```
の権限がLambdaに必要なようです。`AWS管理ポリシーAWSLambdaVPCAccessExecutionRole`に含まれているので、このポリシーをLambdaのロールにアタッチしておきましょう。そうするとVPCが設定できるはずです。

## EFSに接続
用意しておいたEFSとLambdaをつなぎます。コンソールの`ファイルシステム`から`ファイルシステムの追加`をクリックします。EFSファイルシステムとアクセスポイントを指定し、ローカルマウントパスを`/mnt/lambda`として`保存`をクリックします。

## Lambdaでテスト
LambdaからEFSに置いているライブラリを読み込めるかテストしてみます。
EFSは`/mnt/lambda`にマウントされているので、`/mnt/lambda`をpythonのpathに追加しておきます。

```python:lambda_function.py
import json
import sys

sys.path.append("/mnt/lambda")

import torch
import torchvision
import PIL
import cv2
import numpy as np

def lambda_handler(event, context):
    print(f"torch:{torch.__version__}")
    print(f"torchvision:{torchvision.__version__}")
    print(f"PIL:{PIL.__version__}")
    print(f"cv2:{cv2.__version__}")
    print(f"numpy:{np.__version__}")

    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Lambda!')
    }

```

コンソール上で適当なテストイベントを作ってテストを実行すると、下のような結果が得られます。ちゃんとPyTorch読み込めてますね。

```sh
START RequestId: 35329cd4-50f6-4eb7-8950-f27daf75462b Version: $LATEST
OpenBLAS WARNING - could not determine the L2 cache size on this system, assuming 256k
torch:1.6.0+cpu
torchvision:0.7.0+cpu
PIL:7.2.0
cv2:4.4.0
numpy:1.19.2
END RequestId: 35329cd4-50f6-4eb7-8950-f27daf75462b
REPORT RequestId: 35329cd4-50f6-4eb7-8950-f27daf75462b	Duration: 29212.21 ms	Billed Duration: 29300 ms	Memory Size: 128 MB	Max Memory Used: 129 MB	
```

#モデルを動かす
準備は整ったのでモデルを動かしてみます。[EMLのgithubページ](https://github.com/SenJia/EML-NET-Saliency)内に学習済みのモデルの場所が載せてあるのでダウンロードしてきます。READMEの教えに従って、scpなりなんなりでEC2にマウントしておいたEFSの`/mnt/lambda/backbone`に3つのファイル、`res_imagenet.pth, res_places.pth, res_decoder.pth`を置きます。

[eval_combined.py](https://github.com/SenJia/EML-NET-Saliency/blob/master/eval_combined.py)をベースにLambda上で動かせるように修正していきます。

```python:lambda_function.py
import sys
sys.path.append("/mnt/lambda")
import os

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms

from PIL import Image
import cv2
import numpy as np

import resnet
import decoder

# このあたりの環境変数は不要なのでは？という気もしている。
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

image_model_path = "/mnt/lambda/backbone/res_imagenet.pth"
place_model_path = "/mnt/lambda/backbone/res_places.pth"
decoder_model_path = "/mnt/lambda/backbone/res_decoder.pth"
size = (480, 640)
num_feat = 5


def normalize(x):
    x -= x.min()
    x /= x.max()


def post_process(pred):
    pred = cv2.GaussianBlur(pred, (5,5), 10.0)
    normalize(pred)
    pred_uint = (pred * 255).astype(np.uint8)
    return pred, pred_uint


def draw_heatmap(pred, img):
    # saliency mapをもとの画像サイズに変換する。
    resized_pred = np.asarray(Image.fromarray(pred).resize((img.size[0], img.size[1])), dtype=np.uint8)
    resized_colormap = cv2.applyColorMap(resized_pred, cv2.COLORMAP_JET)
    resized_colormap = cv2.cvtColor(resized_colormap, cv2.COLOR_BGR2RGB)
    # 元画像もnumpy ndarrayに変換しておく。
    img_array = np.asarray(img)

    # ブレンディング
    alpha = 0.5
    blended = cv2.addWeighted(img_array, alpha, resized_colormap, 1-alpha, 0)

    return blended


def predict(image_model_path, place_model_path, decoder_model_path, pil_img):
    img_model = resnet.resnet50(image_model_path).eval()
    pla_model = resnet.resnet50(place_model_path).eval()
    decoder_model = decoder.build_decoder(decoder_model_path, size, num_feat, num_feat).eval()

    preprocess = transforms.Compose([
        transforms.Resize(size),
        transforms.ToTensor(),
    ])

    processed = preprocess(pil_img).unsqueeze(0)


    with torch.no_grad():
        img_feat = img_model(processed, decode=True)
        pla_feat = pla_model(processed, decode=True)

        pred = decoder_model([img_feat, pla_feat])

    pred_origin = pred.squeeze().detach().cpu().numpy()
    pred, pred_uint = post_process(pred_origin)

    heatmap = draw_heatmap(pred_uint, pil_img)

    return heatmap


def lambda_handler(event, context):
    # 空のレスポンス
    empty_response = {
        "statusCode": 200,
        "body": "{}"
    }

    pil_img = Image.open("/mnt/lambda/image/examples/115.jpg").convert("RGB")

    heatmap = predict(image_model_path, place_model_path, decoder_model_path, pil_img)
    print(heatmap.shape)

    return empty_response
```
ちょっと長いですが、これでLambda上でSaliencyMapの生成ができます。Lambda上ではGPUが使えないので、元のコードからcuda関連の部分の記述を修正しています。
実行すると、

```
START RequestId: 6f9baccf-b758-4e9a-b43a-b92bdd9757ec Version: $LATEST
OpenBLAS WARNING - could not determine the L2 cache size on this system, assuming 256k
Model loaded /mnt/lambda/backbone/res_imagenet.pth
Model loaded /mnt/lambda/backbone/res_places.pth
Loaded decoder /mnt/lambda/backbone/res_decoder.pth
(511, 681, 3)
END RequestId: 6f9baccf-b758-4e9a-b43a-b92bdd9757ec
REPORT RequestId: 6f9baccf-b758-4e9a-b43a-b92bdd9757ec	Duration: 20075.02 ms	Billed Duration: 20100 ms	Memory Size: 1024 MB	Max Memory Used: 614 MB	
```
となって無事に実行できているようですね。EFS側にファイル出力したかったのですが、若干面倒な感じだったので次回のSlackからの呼び出しのついでにS3に出力できるように修正しましょう。

# まとめ
EFSをLambdaに繋いで、ライブラリの読み込み・モデルファイルの読み込みができるようになりました。これによって従来のようにLambdaの容量を気にせずモデルの推論が実行できるようになるのではないでしょうか。
次回はslack呼び出し編として、今回のLambdaをslackから呼び出して見ようと思います。

# 蛇足
この記事を書くにあたり、一度作成した構成を確認用にもう一度一から作成したのですが、これをやるくらいならAWS CDKとかでズドンとやっておけば執筆が捗ったのになぁと後悔しています。