訓練済み Model の デプロイ

jun-fu@bitengineers です。

5月初旬は日も出て暑かった記憶がありますがその後は雨や曇り続きで、4月の成長から期待していたような変化はないですね。そして記事の内容も盆栽味薄めです。

前回の ML 実践での結果は残念でしたが、せっかくなので AWS SageMaker へデプロイする方法を調べました。以前の記事でもさらっと流れを書いてます。

今回は訓練済みの Model をデプロイする方法を書いてます。主に下記ドキュメントの内容です。

Use PyTorch with the SageMaker Python SDK — sagemaker 2.92.1 documentation

Endpoint 呼び出しの流れ

デプロイ処理を行うと、SageMaker 内に外部からリクエストを受け付ける Endpoint が作成される。作成された Endpoint はリクエストを受け取り、推論を行いレスポンスとして返却する流れとなる。これを可能にするために Endpoint 内部では Model Server が動作しており、事前に訓練済み Model 情報を Load している。この Model Server が Model を Load (読み出し) する処理は model_fn の呼び出しで行われる。

def model_fn(model_dir):
    ...
    return model

この関数で Model が読み出され Model Server が推論可能になる。その後Model Server  がリクエストを受け付けた際の処理の流れとしては以下の 3ステップとなっている。

# Deserialize the Invoke request body into an object we can perform prediction on
input_object = input_fn(request_body, request_content_type)

# Perform prediction on the deserialized object, with the loaded model
prediction = predict_fn(input_object, model)

# Serialize the prediction result into the desired response content type
output = output_fn(prediction, response_content_type)

input_fn

Request をデシリアライズ、するなどし、入力値を用意

predict_fn

入力値と Load された Model を受け取り、推論を行う

output_fn

推論値を Serialize するなどし、Response を作成する

この一連の処理は Model Server で既定の振る舞いが決まっているが、スクリプトで変更することもできる。

この model_fn や input_fn, predict_fn, output_fn を記述するスクリプトをinference script と言う。
model_fn が "must" で実装しなければならないもので、 input_fn, predict_fn, output_fn は任意となる。

デプロイに必要なもの

今回のデプロイは "Endpoint の立ち上げ"とほぼ同意です。Model 情報と model_fn を記述した inference script が必要です。

以降は PyTorch v1.2 or higher を対象として内容です。

Model 情報と書いてますが,これは訓練後の Model を

torch.save(model.state_dict(), "file名")

で書き出したものを使います。ファイル名は ".pt", ".pth" のサフィックスを付けます。PyTorch の略(?)

inference script は任意の名前で用意します。これらのファイルは model.tar.gz という圧縮ファイルで一つにまとめます。 Model Directory Structure に従い、以下の構成にします。

model.pth
code/
  |- inference.py
  |- requirements.txt  # only for versions 1.3.1 and higher

これらを model.tar.gz として S3 の任意の場所にアップロードしておけば後はデプロイするだけです。

実際にデプロイ

まず、私が用意した inference.py は以下のようになってます。

import os
import sys
import copy
import torch
import logging

from bonsai_nn import BonsaiNN

def model_fn(model_dir):
    model = BonsaiNN()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_state_file = os.path.join(model_dir, "model.pth")
    model.load_state_dict(torch.load(model_state_file))
    return model.to(device)

model_fn のみの記述です。bonsai_nn.py から Model class を import しているので code 下に bonsai_nn.py を置いて archive しています。

model.tar.gz/
|- model.pth
|- code/
  |- inference.py
  |- bonsai_nn.py # <= 追加
  |- requirements.txt

この archive を model.tar.gz として S3 に配置します。

訓練済み Model なので PyTorchModel を使い model_data に先ほどS3 に上げた archive を、entry_point に "inference.py" を指定します。

from sagemaker.pytorch.model import PyTorchModel

pytorch_model = PyTorchModel(model_data='s3://be-bonsai/SM_MODEL_DIR/model.tar.gz', role=role,
                             entry_point='inference.py',
                            py_version="py38",
                            framework_version="1.10")

このPyTorchModel インスタンスの deploy() を呼び出してデプロイ完了を待ちます。

predictor = pytorch_model.deploy(instance_type='ml.m5.large', initial_instance_count=1)

しばし待ち、predictor.predict に入力値を渡すと推論が返ってきます。

import numpy as np
import torch

data = [2.49, -1.01, -2.91]
input_data = torch.tensor(np.array(data).astype(np.float32))
print(predictor.predict(input_data)) #[2.49, -1.01, -2.91]) # expected 1.4764e-03

例として、以下が実際に受け取った値です。

[-0.00603854]

これで、Endpoint の動作が確認できました。マネージメントコンソールでも Amazon SageMaker/Endpoints に一覧されています。

使い終わったら

predictor.delete_endpoint()

しないと課金され続けます。

ご注意を。

さいごに

SageMaker・PyTorch に触れてから苦戦してばかりです。
ようやく、外部から推論値が得られるところまではきました。

今回の方法の利点は、

• 採用した Framework(Pytorchなど) での流儀に則って進めたものがそのまま SageMaker へデプロイを試せます

• ローカル環境で Training し期待する結果に見合うか確認した後でデプロイ判断できます

など最初の敷居が低くて良かったです。

欠点は

• Training にローカルリソースが使用される

ので 、理解が進んだ先では Cloud の恩恵を得、 Training も SageMaker でやってしまうと良いのでしょう。

以上です。