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

はじめに

記事「Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する」をもとにSageMakerでElastic Inferenceを使用してPytorchのモデルをホストしようとしたらハマったので、手順を解説します。

Amazon Elastic Inferenceとは

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

PyTorchのサポート

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

参考

• Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する
• オンプレミス環境から Amazon SageMaker を利用する
• Method utils.sts_regional_endpoint(region) always returns the global STS endpoint for all regions?

手順

基本は記事「Amazon Elastic Inference を使用して Amazon SageMaker で PyTorch モデルの ML 推論コストを削減する」の通り進めていきます。

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の認証を設定しておきます。リージョンはモデルをデプロイするリージョンを選びます。以降、使用するリソースは全てここで選んだリージョンで作成する必要があります。

$ 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を使用します。こちらの記事などを読んでJupyter Notebookにアクセスできるようにします。

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

source activate amazonei_pytorch_p36

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

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

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

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/が返ってきてしまい、バージニア北部でしかデプロイできません（参考）。

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

モデルファイル作成

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

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を用意しておきます。デフォルトのハンドラーについてはこちらでプログラムを確認できます。他のコンテナの一覧はこちらから確認できます。
upload_dataではバケットを指定しないので、デフォルトバケットsagemaker-{リージョン名}-{アカウントID}に格納されます。

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分くらいかかります。

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
)

デプロイしたモデルを使用してみます。

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)

エンドポイントの削除

predictor.delete_endpoint()