Microsoft AzureにTorchScriptモデルをデプロイする

公式ドキュメントがごちゃごちゃしていて、個人的にわかりづらかったので、まとめておきます。

公式ドキュメントどんなことが書いてあったか

古いバージョンの内容だったり、sklearnのモデルのデプロイの仕方、azure上でmlflowを用いて実験した際のデプロイの方法が書いてあります。

難しいのは、ばらばらの場所に説明が書いてあったり、古いバージョンでしか動かないコードがそのまま残っている点です。

この記事執筆時点でのバージョンは、

azureml-core==1.17.0

です。

pytorchモデルはすでに torchscript 化している前提で進めます。

下準備

パッケージ

pip install azureml-core

まだ入ってなければいれてください

以下、notebookの使用を想定

パラメータ

使うパラメータを定義しておきます(ご自身のパラメータを設定してください)。

WORKSPACE_NAME = "ws_name"
WORKSPACE_LOCATION = "East US"
SUBSCRIPTION_ID = ""
RESOURCE_GROUP = ""

MODEL_NAME = "model_name"
MODEL_DESCRIPTION = """ model description """
ENVIRONMENT_NAME = "env_name"
SERVICE_NAME = "torchscript-service"

TORCHSCRIPT_PATH = "./trace_model.zip"

インポート

デプロイするだけならtorchはインポートしなくても大丈夫です。

from datetime import datetime
import json

from azureml.core import Model
from azureml.core import Workspace
from azureml.core import Environment
from azureml.core.conda_dependencies import CondaDependencies
from azureml.core.webservice import AciWebservice
from azureml.core.model import InferenceConfig

azureml をごちゃごちゃ

workspaceの作成

Machine Learning ワークスペースが必要です。以下のコマンドではワークスペースがなければ作成され、あればそれが利用されます。必要ならブラウザに飛んで認証します。手動でやるときは忘れがちです。

ws = Workspace.create(name = WORKSPACE_NAME,
                      location=WORKSPACE_LOCATION,
                      resource_group=RESOURCE_GROUP,
                      subscription_id=SUBSCRIPTION_ID,
                      exist_ok=True
                     )
print(f"workspace name: {ws.name}")

モデルのregister

「登録」とは言いますが、実質的にはファイルかディレクトリをアップロードして名前を付けて紐づけているだけになります

your_model = Model.register(
  workspace=ws,
  model_path=TORCHSCRIPT_PATH,
  model_name=MODEL_NAME,
  model_framework=Model.Framework.CUSTOM,
  description=MODEL_DESCRIPTION
)

環境

azure container imageにデプロイする際の、コンテナの環境を設定します。今度はtorchscriptをロードして実行するためにtorchが必要になります。といっても、必要である旨を教えてあげればいいです。

たとえば以下のようになります。

environment = Environment(ENVIRONMENT_NAME)
environment.python.conda_dependencies = CondaDependencies.create(pip_packages=[
    'azureml-defaults',
    'inference-schema[numpy-support]',
    'numpy',
    'torch'
])

score.pyを書く

ここまではnotebook上で実行してきましたが、container上でどういった処理を行うかという部分はpythonファイルに書いておく必要があります(ファイルと一緒にcontainerに送ります)。

処理といっても、リクエストを受け取ったあとにそれをどうパースしてモデルに入力して出力を得るかの部分だけで、わかりやすいです。また、それなりに自由度もあるので、扱いやすいです。ドキュメントさえ何とかしてくれればいいのに

要件としては、init()とrun(data)を実装する必要があります

import os
import json
import time
import torch

TORCH_SCRIPT_FILENAME = "./trace_model.zip"  # *change here

def init():
    global model
    
    model_path = os.path.join(os.getenv("AZUREML_MODEL_DIR"), TORCH_SCRIPT_FILENAME)
    model = torch.jit.load(model_path)

def run(data):
    try:
  data = json.loads(data)
       input = do_sth(data["data"])
  output = model(input)
        return {
            "any": output,
        }
    except Exception as e:
        print(e)
        return e

initのAZUREML_MODEL_DIRが、アップロードしたファイル等が置かれているパスという認識です。

modelはglobalで定義します。後でrun内で使うためです

リクエストの"data"に送ったデータは入ります。

responseは好きに書けばいいでしょう。

デプロイ

deployの設定

notebookに戻ります。

aci_config = AciWebservice.deploy_configuration(cpu_cores=1, memory_gb=2)

inference_config = InferenceConfig(
    entry_script='score.py', 
    environment=environment
)

service = Model.deploy(
    workspace=ws,
    name=SERVICE_NAME,
    models=[your_model],
    inference_config=inference_config,
    deployment_config=aci_config,
    overwrite=True
)

service.wait_for_deployment(show_output=True)

これで待てばデプロイされます。長くて6分くらいです。