SageMakerでElastic Inferenceを使用してPytorchのモデルをホストする

# はじめに

記事「[Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する](https://aws.amazon.com/jp/blogs/news/reduce-ml-inference-costs-on-amazon-sagemaker-for-pytorch-models-using-amazon-elastic-inference/)」をもとにSageMakerでElastic Inferenceを使用してPytorchのモデルをホストしようとしたらハマったので、手順を解説します。

### Amazon Elastic Inferenceとは
GPUメモリが小さいGPUをEC2やECSタスクにアタッチすることで、機械学習の推論のホストを効率的にするサービスです。トレーニング時は必要だったメモリは、推論のホスト時はそこまで必要でないケースがほとんどです。GPUメモリは2GB, 4GB, 8GBと小さい代わりに、[コストを削減](https://aws.amazon.com/jp/machine-learning/elastic-inference/pricing/)することができます。基本はEC2へアタッチして使用しますが、ECSのタスクにもアタッチできるので、コンテナで複数ジョブを走らせてスケールさせることもでき（ると思い）ます。



### PyTorchのサポート
Elastic Inferenceに対応しているPytorchを使用する必要があるので、特別な環境を構築する必要があります。2020/08時点でSageMakerのノートブックでサポートしていないので、EC2を使用してモデルを構築、SageMakerにデプロイする必要があります。

# 参考
- [Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する](https://aws.amazon.com/jp/blogs/news/reduce-ml-inference-costs-on-amazon-sagemaker-for-pytorch-models-using-amazon-elastic-inference/)
- [オンプレミス環境から Amazon SageMaker を利用する](https://aws.amazon.com/jp/blogs/news/sagemaker_from_onpremises/)
- [Method utils.sts_regional_endpoint(region) always returns the global STS endpoint for all regions?](https://www.gitmemory.com/issue/aws/sagemaker-python-sdk/1332/596423188)

# 手順

基本は記事「[Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する](https://aws.amazon.com/jp/blogs/news/reduce-ml-inference-costs-on-amazon-sagemaker-for-pytorch-models-using-amazon-elastic-inference/)」の通り進めていきます。

## IAMユーザー作成
`ElasticInferenceBuild`という名前でユーザーを作成しました。アクセスの種類は`プログラムによるアクセス`、ポリシーは`AmazonSageMakerFullAccess`をアタッチしました。
このユーザーはSageMakerを操作するために、EC2上のプログラムで使用します。

## IAMロール作成
`ElasticInferenceBuild`という名前でロールを作成しました。信頼されたエンティティの種類は`SageMaker`にし、ポリシーは`AmazonSageMakerFullAccess`、`AmazonS3FullAccess`、`AmazonEC2ContainerRegistryFullAccess`をアタッチしました。
このロールはSageMakerにassumeし、SageMakerがS3やECRにアクセスできるようにするために使用します。

## EC2インスタンス作成
AMIはAWS Marketplaceで`DLAMI`で検索すると出てくる`Deep Learning AMI (Amazon Linux) V32.0`でインスタンスを作成します。後に使用する`amazonei_pytorch_p36`というElasic Inference対応の仮想環境が入っているのは現状このインスタンスのみのようです。

モデル構築にはメモリが必要になるので、メモリが大きい`c5.xlarge`あたりを選んでおくとよいです。

sshとhttp（Jupyter Notebook）アクセスができるようにセキュリティグループを設定しておきます。

## 認証情報の設定

EC2にログインしたら、ユーザー`ElasticInferenceBuild`の認証を設定しておきます。リージョンはモデルをデプロイするリージョンを選びます。以降、使用するリソースは全てここで選んだリージョンで作成する必要があります。

```bash
$ aws configure
AWS Access Key ID [None]: AKIAIOSFODNN7EXAMPLE
AWS Secret Access Key [None]: wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY
Default region name [None]: ap-northeast-1
Default output format [None]: json
```

## ライブラリ更新
スクリプトでもいいのですが、操作しやすいのでJupyter Notebookを使用します。[こちらの記事](https://qiita.com/taka4sato/items/2c3397ff34c440044978)などを読んでJupyter Notebookにアクセスできるようにします。

condaの仮想環境で`amazonei_pytorch_p36`をアクティベートします。

```bash
source activate amazonei_pytorch_p36
```

ワーク用のディレクトリを用意して、notebookを作成します。
まず、SageMakerとbotocoreのSDKが古いので、これらをアップデートします。ターミナルではうまくいかない場合があるので、notebook上で実行します。

```python
!pip install botocore==1.17.44
!pip install sagemaker==2.4
```

インストールできたか確認します。

```python
import botocore
import sagemaker

print('botocore:', botocore.__version__)
print('sagemaker:', sagemaker.__version__)
print(sagemaker.utils.sts_regional_endpoint('ap-northeast-1'))

### 出力
# botocore: 1.17.44
# sagemaker: 2.4.0
# https://sts.ap-northeast-1.amazonaws.com
```

これをしないと`sagemaker.utils.sts_regional_endpoint('ap-northeast-1')`が、どのリージョンを指定してもグローバルエンドポイント（バージニア北部）`http://sts.amazonaws.com/`が返ってきてしまい、バージニア北部でしかデプロイできません（[参考](https://www.gitmemory.com/issue/aws/sagemaker-python-sdk/1332/596423188)）。

また、SageMaker SDKがバージョン2系になるとインターフェースが微妙に変わるので、以降のプログラムでは書き換えが必要です。


## モデルファイル作成
モデルをTorchScriptに変換することが必要です。今回は`torchvision`の`DenseNet121`を例にモデルを作成します。
Elastic Inference対応の`torch`が入っているので`torch.jit.optimized_execution`の引数が２つ似変更されています。Elastic Inferenceの序数を入れます。
importでクラッシュする場合は、`sagemaker`→`torch`の順で読み込むと大丈夫な場合があります。

```python
import sagemaker
from sagemaker.pytorch import PyTorchModel
import torch, torchvision
import subprocess


# Toggle inference mode
model = torchvision.models.densenet121(pretrained=True).eval()
cv_input = torch.rand(1,3,224,224)

# Required when using Elastic Inference
with torch.jit.optimized_execution(True, {'target_device': 'eia:0'}):
    model = torch.jit.trace(model, cv_input)

model = torch.jit.trace(model, cv_input)
torch.jit.save(model, 'model.pt')
subprocess.call(['tar', '-czvf', 'densenet121_traced.tar.gz', 'model.pt'])
```
モデルはtarballにします。これをS3に格納します。

`region`にはデプロイしたいリージョン、`role`には`ElasticInferenceBuild`のARNを指定します。
`instance_type`にはCPUインスタンスの指定を、`accelerator_type`にはアタッチするElastic Inferenceを指定します。
コンテナはElastic Inference対応のPytorchがインストールされているコンテナを使用します。現状PyTorch1.3.1しかないようです。また、ECRはリージョンごとに用意されているので、デプロイするリージョンに合わせます。コンテナはデフォルトで`model_fn`などを持っているので、エントリポイントは空のファイル`script.py`を用意しておきます。デフォルトのハンドラーについては[こちら](https://github.com/aws/sagemaker-pytorch-inference-toolkit/blob/master/src/sagemaker_pytorch_serving_container/default_pytorch_inference_handler.py)でプログラムを確認できます。他のコンテナの一覧は[こちら](https://docs.aws.amazon.com/deep-learning-containers/latest/devguide/deep-learning-containers-images.html)から確認できます。
`upload_data`ではバケットを指定しないので、デフォルトバケット`sagemaker-{リージョン名}-{アカウントID}`に格納されます。

```python
sagemaker_session = sagemaker.Session()
region = 'ap-northeast-1'
role = 'arn:aws:iam::xxxxxxxxxxxx:role/ElasticInferenceBuild'

instance_type = 'c5.large'
accelerator_type = 'eia2.medium'

ecr_image = '763104351884.dkr.ecr.{}.amazonaws.com/pytorch-inference-eia:1.3.1-cpu-py3'.format(region)

# Satisfy regex
endpoint_name = 'pt-ei-densenet121-traced-{}-{}'.format(instance_type, accelerator_type).replace('.', '').replace('_', '')
tar_filename = 'densenet121_traced.tar.gz'

# script.py should be blank to use default EI model_fn and predict_fn
# For non-EI PyTorch usage, must implement own model_fn
entry_point = 'script.py'

# Returns S3 bucket URL
model_data = sagemaker_session.upload_data(path=tar_filename)
print('Uploadded tarball to', model_data)

### 出力
# Uploadded tarball to s3://sagemaker-ap-northeast-1-xxxxxxxxxxxx/data/densenet121_traced.tar.gz
```

次のコマンドでエンドポイントをデプロイします。コンソールでエンドポイントのデプロイを確認できます。10分くらいかかります。

```python
pytorch = PyTorchModel(
    model_data=model_data,
    framework_version='1.3.1',
    py_version='py36',
    role=role,
    image_uri=ecr_image,
    entry_point=entry_point,
    sagemaker_session=sagemaker_session
)

# Function will exit before endpoint is finished creating
predictor = pytorch.deploy(
    initial_instance_count=1,
    instance_type='ml.' + instance_type,
    accelerator_type='ml.' + accelerator_type,
    endpoint_name=endpoint_name,
    wait=False
)
```
デプロイしたモデルを使用してみます。

```python
from sagemaker.pytorch import PyTorchPredictor


predictor = PyTorchPredictor(endpoint_name)
data = torch.rand(1,3,224,224)
output = predictor.predict(data)

print(type(output))
print(output.shape)

### 出力
# <class 'numpy.ndarray'>
# (1, 1000)
```

## エンドポイントの削除

```python
predictor.delete_endpoint()
```