TorchServe を使用した大規模な推論のための PyTorch モデルをデプロイする

今日ご紹介するサービスの多くは機械学習 (ML) を必要とするものです。オンライン検索や製品のレコメンデーションから音声認識や翻訳まで、これらのサービスには予測を提供するための ML モデルが必要です。ML がさらに多くのサービスで採用されるにつれて、苦労して結果を得て、モデルを迅速かつ確実に本番環境にデプロイすることが困難になります。そして、これらのサービスを利用する人の数が増えるにつれて、モデルが数百万のユーザーに同時に、そして、安全かつ確実に、低レイテンシーで予測を提供できるようにすることには、より大きな困難が伴います。

開発者は、モデル開発にさまざまなオープンソースフレームワークを使用します。ここ数年、PyTorch は、ML を利用したアプリケーションを開発する多くの研究者、開発者、およびデータサイエンティストが選択する深層学習のフレームワークとなっています。PyTorch は、そのシンプルさに加えて、Python 的な方法でモデルを実装およびトレーニングでき、Eager モードと Graph モードをシームレスに切り替える機能があることから好まれています。しかしながら、これまで、PyTorch モデルを本番環境で大規模に提供することについて、簡単かつネイティブにサポートされた方法はありませんでした。

AWS は、PyTorch のオープンソースモデルサービスライブラリである TorchServe の実験的リリースを共有できることを嬉しく思います。

AWS は、Facebook と連携して TorchServe を開発しました。AWS と Facebook は、より広範な PyTorch コミュニティと共に、TorchServe に注力し続けます。クラウドベースの PyTorch プロジェクトが 83% を超えて AWS で行われていることから、PyTorch モデルのデプロイの困難に対処するために TorchServe を立ち上げることができることを幸いに思います。TorchServe を使用すると、PyTorch モデルを TorchScript を使用して Eager モードまたは Graph モードでデプロイし、複数のモデルを同時に提供し、A/B テスト用に本番モデルをバージョン管理し、モデルを動的にロードおよびアンロードし、詳細なログとカスタマイズ可能なメトリクスを監視できます。

TorchServe は使いやすいです。ローカルにデプロイするのに便利な CLI が付属しており、コンテナにパッケージ化して Amazon SageMaker または Amazon EKS で簡単にスケールアウトできます。画像分類、オブジェクト検出、画像セグメンテーション、テキスト分類などの一般的な問題のデフォルトハンドラーを使用して、数行のコードでデプロイできます。初期化、前処理、後処理用の長いサービスハンドラーを記述する必要はありません。TorchServe はオープンソースです。つまり、完全にオープンであり、デプロイのニーズに合わせて拡張できます。

この投稿では、TorchServe、その主な特徴と機能、およびその使用方法について詳しく説明します。サンプルコードを提供して、利点と主要な概念を説明し、TorchServe と Amazon SageMaker を使用して PyTorch 推論をスケーリングする例も示します。

TorchServe のリリース前は、PyTorch モデルを提供する場合、独自のモデル提供ソリューションを開発する必要がありました。モデルのためのカスタムハンドラーを開発し、モデルサーバーを開発し、独自の Docker コンテナを構築し、ネットワーク経由でアクセスできるようにする方法を理解し、クラスターオーケストレーションシステムと統合するなど、多くのことをする必要がありました。

TorchServe を使用すると、多くの機能をすぐに使用できます。人気のあるモデルのためのカスタムハンドラーを記述する必要なく、トレーニング済みの PyTorch モデルを大規模に高いパフォーマンスでデプロイする完全な柔軟性を提供します。ほんの数行のコードで、トレーニング済みのモデルを本番環境にデプロイできます。

TorchServe の開始方法

TorchServe の使用は、簡単に開始できます。この投稿では、Deep Learning AMI を実行する c5.xlarge Amazon EC2 インスタンスでサンプルをテストしました。ローカルのノートパソコンまたはデスクトップで手順を試すこともできます。Amazon EC2 インスタンスを起動する手順については、Amazon EC2 の開始方法を参照してください。

TorchServe をインストールするには、GitHub の手順に従ってください。依存関係を管理するには、Conda または他の仮想環境を使用することをお勧めします。TorchServe をインストールすると、次の手順を実行して最初のモデルをデプロイできます。

1. サンプルにアクセスするには、TorchServe リポジトリをダウンロードしてください。次のコードを実行します。

   mkdir torchserve-examples
   cd torchserve-examples
   git clone https://github.com/pytorch/serve.git

2. DenseNet の画像分類モデルを公式の PyTorch モデルリポジトリからダウンロードします。次のコードを実行します。

   wget https://download.pytorch.org/models/densenet161-8d451a50.pth

3. モデルを PyTorch から TorchServe 形式に変換します。TorchServe は、拡張子 .mar のモデルアーカイブ形式を使用します。.mar ファイルは、state_dict (各レイヤーをそのパラメータテンソルにマッピングする辞書オブジェクト) を使用して、モデルチェックポイントまたはモデル定義ファイルをパッケージ化します。TorchServe の torch-model-archiver ツールを使用して、.mar ファイルを作成できます。カスタムハンドラーを作成する必要はありません。--handler image_classifier を指定するだけで、自動的にハンドラーがセットアップされます。これで .mar ファイルができたので、TorchServe でホストします。次のコードを実行します。

   torch-model-archiver --model-name densenet161 \
   --version 1.0 --model-file serve/examples/image_classifier/densenet_161/model.py \
   --serialized-file densenet161-8d451a50.pth \
   --extra-files serve/examples/image_classifier/index_to_name.json \
   --handler image_classifier
   
   ls *.mar
   

   次の出力が得られます。

   densenet161.mar serve

4. モデルをホストします。次のコードを実行します。

   mkdir model_store
   mv densenet161.mar model_store/
   torchserve --start --model-store model_store --models densenet161=densenet161.mar

5. 同じホストで別の端末を開き、次のコードを実行して、TorchServe をテストします (tmux を使用して複数のセッションを管理できます)。

   curl -O https://s3.amazonaws.com/model-server/inputs/kitten.jpg
   curl -X POST http://127.0.0.1:8080/predictions/densenet161 -T kitten.jpg

   次の出力が得られます。

   [
    {
    "tiger_cat": 0.4693356156349182
    },
    {
    "tabby": 0.46338796615600586
    },
    {
    "Egyptian_cat": 0.06456131488084793
    },
    {
    "lynx": 0.0012828155886381865
    },
    {
    "plastic_bag": 0.00023323005007114261
    }
   ]

デフォルトハンドラーを備えた標準モデルのゼロコード変更デプロイ

多くの深層学習のユースケースは、画像分類、オブジェクト検出、画像セグメンテーション、テキスト分類のいずれかに分類されます。これらのアプリケーションのいずれかを使用している場合は、前のセクションのように、コードを変更せずにデプロイできます。Eager から TorchScript に、またはその逆に変換したり、初期化、前処理、および後処理のためのサービスハンドラーを記述したりする必要はありません。

TorchServe の torch-model-archiver ツールは、PyTorch のさまざまな表現 (Eager モードと TorchScript) を自動的に検出して処理できます。TorchVision、TorchText、TorchAudio などのパッケージでサポートされる一般的なモデルの場合、torch-model-archiver は初期化、前処理、および後処理にデフォルトハンドラーを使用します。カスタムモデルがある場合や、デフォルトハンドラーを拡張するカスタムロジックを導入する場合でも、カスタムハンドラーを記述できます。

サポートされているデフォルトハンドラーの完全なリストについては、GitHub リポジトリを参照してください。独自のカスタムハンドラーを記述する手順については、GitHub リポジトリを参照してください。

複数のモデルのホスティングとワーカーのスケーリング

TorchServe は、登録済みモデルの一覧表示、既存のサーバーへの新しいモデルの登録、現在のモデルの登録解除、モデルあたりのワーカー数の増減、モデルのステータスの説明、バージョンの追加、デフォルトバージョンの設定を行う管理 API を提供します。管理 API はポート 8081 でリッスンしており、デフォルトではローカルホストからのみアクセスできますが、デフォルトの動作を変更できます。

新しいモデルを登録するには、以下の手順を実行します。

1. 次のコードで新しいモデルをダウンロードします。

   wget https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
   torch-model-archiver --model-name fastrcnn --version 1.0 \
   --model-file serve/examples/object_detector/fast-rcnn/model.py \
   --serialized-file fasterrcnn_resnet50_fpn_coco-258fb6c6.pth \
   --handler object_detector \
   --extra-files serve/examples/object_detector/index_to_name.json
   
   mv fastrcnn.mar model_store/

2. 新しいモデルを次のコードで登録します。

   curl -X POST "http://localhost:8081/models?url=fastrcnn.mar"
   

   次の出力が得られます。

   {
    "status": "Model \"fastrcnn\" registered"
   }
   

   次のコードを使用して、登録済みモデルのリストをクエリすることもできます。

   curl "http://localhost:8081/models"

   次の出力が得られます。

   {
       "models": [
           {
               "modelName": "densenet161",
               "modelUrl": "densenet161.mar"
           },
           {
               "modelName": "fastrcnn",
               "modelUrl": "fastrcnn.mar"
           }
       ]
   }
   

3. モデルのワーカーをスケーリングします。新しいモデルにはワーカーが割り当てられていないため、次のコードでワーカーの最小数を設定します。

   curl -v -X PUT "http://localhost:8081/models/fastrcnn?min_worker=2"
   curl "http://localhost:8081/models/fastrcnn"
   

   次の出力が得られます。

   [
     {
       "modelName": "fastrcnn",
       "modelVersion": "1.0",
       "modelUrl": "fastrcnn.mar",
       "runtime": "python",
       "minWorkers": 2,
       "maxWorkers": 2,
       "batchSize": 1,
       "maxBatchDelay": 100,
   …
   …

   モデルが、96 個の vCPU を搭載した c5.24xlarge EC2 インスタンスなど、多くのコアを搭載した CPU でホストされている場合、前述の方法を使用して、スレッド数を簡単にスケーリングできます。

4. 次のコードでモデルの登録を解除します。

   curl -X DELETE http://localhost:8081/models/fastrcnn/

5. モデルをバージョン管理するには、torch-model-archiver を呼び出すときに、バージョン番号を --version に渡します。次のコードを参照してください。

   torch-model-archiver --model-name fastrcnn --version 1.0 ...

バッチ予測 (またはバッチ推論) の実行

アプリケーションによっては、推論をバッチで実行する必要がある場合があります。大規模なデータセットがあり、オフラインで推論を生成したい場合は、データセットを大きなバッチに集約して処理する方が計算上効率的です。レイテンシーの許容度が高い一部のリアルタイムユースケースでは、複数のリクエストをバッチ処理して結果を提供し、リソース使用率を改善して、プロセスの運用コストを削減できます。バッチ処理のもう 1 つのユースケースは、新しいモデルをトレーニングする前に、トレーニングデータを他のモデルで前処理することです。

TorchServe は、バッチ推論をネイティブにサポートしています。TorchServe バッチ推論を実行する手順については、GitHub リポジトリを参照してください。

ロギングとメトリクス

TorchServe を使用すると、完全にカスタマイズ可能なログやメトリクスへのアクセスが容易になります。デフォルトでは、TorchServe はログメッセージを stderr および stout に出力します。TorchServe は log4j を使用し、log4j プロパティを変更することでロギングをカスタマイズできます。

TorchServe は、デフォルトで CPUUtilization、DiskUtilization などのシステムレベルのメトリクスも収集します。メトリクス API を使用してカスタムメトリクスを指定することもできます。次のスクリーンショットは、TorchServe から推論がリクエストされたときのデフォルトのログ出力を示しています。

TorchServe と Amazon SageMaker を使用した大規模な PyTorch デプロイ

Amazon SageMaker は、データのラベル付け、モデル開発、大規模なモデルのトレーニング、モデルのデプロイのための機能を備えたフルマネージド型のサービスです。デプロイとホスティングについて、Amazon SageMaker は、Amazon SageMaker トレーニングクラスターでトレーニングされたモデルの便利なワンクリックデプロイを提供します。また、完全にモジュール化されているため、独自のアルゴリズムとコンテナを取り込んで、必要なサービスのみを使用できます。

この投稿では、TorchServe コンテナを構築し、Amazon SageMaker を使用してホストする方法を示します。Amazon SageMaker は、フルマネージド型のホスティング体験を提供します。インスタンスのタイプと、希望する最大数および最小数を指定するだけです。他の作業は、Amazon SageMaker が行います。

数行のコードで、インスタンスの起動、Amazon S3 から TorchServe コンテナへのモデルのダウンロード、アプリケーション用の安全な HTTPS エンドポイントのセットアップを Amazon SageMaker に指示できます。クライアント側では、TorchServe がサポートするこの安全なエンドポイントへの簡単な API 呼び出しで予測を取得します。

サンプルの実行

このサンプルは、Amazon SageMaker Notebook インスタンス、Amazon EC2、またはノートパソコンやデスクトップで実行できます。ローカルのノートパソコンまたはデスクトップを使用している場合は、AWS CLI をダウンロードしてインストールし、それとともに、AWS SDK for Python (boto3) と Amazon SageMaker Python SDK を設定します。デプロイ後、モデルは Amazon SageMaker のフルマネージド型のデプロイインスタンスでホストされます。

最も便利にするには、Jupyter ノートブックインターフェイスを提供し、すべての AWS ライブラリがインストールされており、すぐに使用できる Amazon SageMaker ノートブックインスタンスを起動します。詳細については、Amazon SageMaker ノートブックインスタンスの作成を参照してください。

この投稿は、Amazon SageMaker ノートブックインスタンスで次の手順を実行していることを前提としています。

IAM のロールとポリシー

新しいノートブックインスタンスを作成すると、新しい IAM ロールを作成するように求められます。Amazon S3 アクセスを持つロールを作成しようとしています。Amazon SageMaker コンソールは、このプロセスをガイドします。

新しいロールは XXX の一意の識別子を持つ AmazonSageMaker-ExecutionRole-XXX のようになります。TorchServe コンテナをビルドして Amazon ECR にプッシュする必要があるため、AmazonEC2ContainerRegistryFullAccess ポリシーをノートブックインスタンスロールに追加する必要もあります。これは、IAM コンソールで、またはAWS CLI を使用して実行できます。次のコードを参照してください。

aws iam attach-role-policy \
    --policy-arn arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryFullAccess \
    --role-name AmazonSageMaker-ExecutionRole-XXX

詳細は、IAM ID のアクセス許可の追加および削除を参照してください。

実践

この投稿で使用されているコード、設定ファイル、Jupyter ノートブック、Dockerfiles は GitHub で入手できます。次のサンプルの手順は、deploy_torchserve.ipynb Jupyter ノートブックからのものです。

実践のため、deploy_torchserve.ipynb ノートブックを開き、各セルを実行します。または、次の手順をコピーして新しい Jupyter ノートブックの別のセルに貼り付け、順番に実行することもできます。

次のスクリーンショットは、deploy_torchserve.ipynb のサンプルを示しています。

サンプルリポジトリの複製

サンプルリポジトリのクローンを作成するには、次のコードを入力します。

git clone https://github.com/shashankprasanna/torchserve-examples.git
cd torchserve-examples

TorchServe リポジトリのクローンを作成し、torch-model-archiver をインストールします

「torch-model-archiver」ツールを使用して、モデルアーカイブファイルを作成します。.mar モデルアーカイブファイルには、モデルのチェックポイントとその「state_dict」(各レイヤーをそのパラメータテンソルにマッピングする辞書オブジェクト) が含まれています。

!git clone https://github.com/pytorch/serve.git
!pip install serve/model-archiver/

PyTorch モデルのダウンロード、TorchServe アーカイブの作成、および Amazon S3 へのアップロード

PyTorch モデルをダウンロードして TorchServe アーカイブを作成するには、次のコードを入力します。

!wget -q https://download.pytorch.org/models/densenet161-8d451a50.pth

model_file_name = 'densenet161'
!torch-model-archiver --model-name {model_file_name} \
--version 1.0 --model-file serve/examples/image_classifier/densenet_161/model.py \
--serialized-file densenet161-8d451a50.pth \
--extra-files serve/examples/image_classifier/index_to_name.json \
--handler image_classifier

モデルの Amazon S3 へのアップロード

モデルを Amazon S3 にアップロードするには、次の手順を実行します。

1. boto3 セッションを作成し、リージョンとアカウント情報を取得します

   import boto3, time, json
   sess    = boto3.Session()
   sm      = sess.client('sagemaker')
   region  = sess.region_name
   account = boto3.client('sts').get_caller_identity().get('Account')
   
   import sagemaker
   role = sagemaker.get_execution_role()
   sagemaker_session = sagemaker.Session(boto_session=sess)
   

2. デフォルトの Amazon SageMaker S3 バケット名を取得します

   bucket_name = sagemaker_session.default_bucket()
   prefix = 'torchserve'
   

3. Amazon SageMaker はモデルが tar.gz ファイルにあると想定しているため、densenet161.mar ファイルから圧縮 tar.gz ファイルを作成します。

   !tar cvfz {model_file_name}.tar.gz densenet161.mar

4. モデルのディレクトリの下の S3 バケットにモデルをアップロードします。

!aws s3 cp {model_file_name}.tar.gz s3://{bucket_name}/{prefix}/model

Amazon ECR レジストリの作成

TorchServe コンテナイメージ用の新しい Docker コンテナレジストリを作成します。Amazon SageMaker は、このレジストリから TorchServe コンテナをプルします。次のコードを参照してください。

registry_name = 'torchserve'
!aws ecr create-repository --repository-name torchserve

TorchServe Docker コンテナの作成および Amazon ECR へのプッシュ

この投稿のリポジトリには、TorchServe コンテナを構築するための Dockerfile がすでに含まれています。Docker コンテナイメージをローカルでビルドし、前のステップで作成した Amazon ECR リポジトリにプッシュします。次のコードを参照してください。

image_label = 'v1'
image = f'{account}.dkr.ecr.{region}.amazonaws.com/{registry_name}:{image_label}'

!docker build -t {registry_name}:{image_label} .
!$(aws ecr get-login --no-include-email --region {region})
!docker tag {registry_name}:{image_label} {image}
!docker push {image}

コンテナが構築され、Amazon ECR に正常にプッシュされたことを確認する次の出力が表示されます。

推論エンドポイントのホスト

推論エンドポイントをホストして予測を行う方法は複数あります。最も迅速なアプローチは、Amazon SageMaker Python SDK を使用することです。ただし、クライアントアプリケーションからエンドポイントを呼び出す場合は、選択した言語で Amazon SDK を使用する必要があります。

推論エンドポイントのホストおよび Amazon SageMaker SDK での予測

推論エンドポイントをホストし、Amazon SageMaker SDK を使用して予測を行うには、次の手順を実行します。

1. モデルを作成します。モデル関数は、TorchServe コンテナイメージの名前とトレーニング済みモデルのロケーションを想定しています。次のコードを参照してください。

   import sagemaker
   from sagemaker.model import Model
   from sagemaker.predictor import RealTimePredictor
   role = sagemaker.get_execution_role()
   
   model_data = f's3://{bucket_name}/models/{model_file_name}.tar.gz'
   sm_model_name = 'torchserve-densenet161'
   
   torchserve_model = Model(model_data = model_data,
                           image = image,
                           role = role,
                           predictor_cls=RealTimePredictor,
                           name = sm_model_name)
   

   モデル関数の詳細については、モデルを参照してください。

2. Amazon SageMaker コンソールでモデルの詳細を表示するには、[モデル] を選択します。
   
3. モデルエンドポイントをデプロイします。Amazon SageMaker でコンテナを実行するインスタンスタイプとインスタンスの数を指定します。次のコードを参照してください。

   endpoint_name = 'torchserve-endpoint-' + time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime())
   predictor = torchserve_model.deploy(instance_type='ml.m4.xlarge',
    initial_instance_count=1,
    endpoint_name = endpoint_name)
   

   呼び出しの合計数などのメトリクスに基づいて自動的にスケーリングするようにセットアップすることもできます。詳細については、Amazon SageMaker モデルの自動スケーリングを参照してください。

4. Amazon SageMaker コンソールで、ホストされているエンドポイントを表示するには、[エンドポイント] を選択します。
   
5. 次のコードでモデルをテストします。

   !wget -q https://s3.amazonaws.com/model-server/inputs/kitten.jpg 
   file_name = 'kitten.jpg'
   with open(file_name, 'rb') as f:
    payload = f.read()
    payload = payload
   
   response = predictor.predict(data=payload)
   print(*json.loads(response), sep = '\n')
   

次のスクリーンショットは、TorchServe によってホストされているモデルの呼び出しの出力を示しています。モデルは、画像の子猫はジャガーネコまたはトラネコのいずれかであると考えています。

新しいデータの予測を取得するために TorchServe エンドポイントを呼び出す必要があるモバイルアプリやウェブページなどのアプリケーションを構築している場合は、Amazon SageMaker SDK ではなく Amazon API を使用できます。たとえば、クライアント側で Python を使用している場合は、Amazon SDK for Python (boto3) を使用します。boto3 を使用してモデルを作成し、エンドポイントを設定し、エンドポイントを作成し、最後に推論エンドポイントで推論を実行する方法の例については、GitHub のこのサンプルの Jupyter ノートブックを参照してください。

まとめ

この投稿では、TorchServe とその主な機能および利点を紹介しました。TorchServe は、モデルを本番環境で使用できるようにする開発者と、本番環境でコンテナをデプロイする Ops エンジニアの両方にとって使いやすいものです。TorchServe は、Eager モードと TorchScript をサポートし、最も一般的にデプロイされるモデルのデフォルトハンドラーが付属しているため、コード変更を要することなくデプロイできます。TorchServe は複数のモデルを同時にホストでき、バージョン管理をサポートします。機能の完全なリストについては、GitHub リポジトリを参照してください。

この投稿では、Amazon SageMaker を使用して TorchServe に PyTorch モデルをデプロイするエンドツーエンドのデモも紹介しました。これをテンプレートとして使用して、独自の PyTorch モデルを Amazon SageMaker にデプロイできます。完全なサンプルは、GitHub で参照できます。

TorchServe またはこの投稿について質問やコメントがある場合は、以下でコメントをお寄せいただくか、GitHub リポジトリでイシューを作成してください。

著者について

Shashank Prasanna 氏は、アマゾン ウェブ サービス (AWS) の AI および機械学習テクニカルエバンジェリストです。機械学習に関する難題を解決するエンジニア、開発者、データサイエンティストのサポートに注力しています。AWS に入社する前は、NVIDIA、MathWorks (MATLAB および Simulink のメーカー)、Oracle で、プロダクトマーケティング、プロダクトマネジメント、ソフトウェア開発に携わっていました。

Manoj Rao 氏は、AWS のソフトウェア開発者です。ML Runtime/Inference の AWS が誇る品質を、さらに可能な限り高めるように尽力しています。時間があるときには、ブログで Linux Kernel、Emacs などについて投稿しています。