こんにちは。 技術本部研究開発部の高橋寛治です。
SageMakerのInf1推論エンドポイントで機械学習モデルを試したので紹介します。 部分的に紹介しているため、SageMakerやHuggingFaceを使ったことがないと、わかりづらいかと思いますがご了承ください。
目的
現在SageMakerの推論エンドポイントで、独自コンテナの機械学習モデルを稼働させています。 これをコスパよく、より高速に動かしたいというのが今回の目的です。
Inf1インスタンス
Inf1インスタンスは、機械学習の推論に特化したAWS Inferentiaチップを搭載したインスタンスのことです。 GPUインスタンスよりもスループットが高く、推論あたりのコストが低いことが特徴です。
Inf1インスタンスで上記の恩恵にあずかるために、「モデルのコンパイル」と「推論環境の構築」が必要となります。
モデルのコンパイル
今回コンパイルするモデルは、HuggingFaceのTransformersライブラリを利用した固有表現抽出モデルです。
環境構築は、AWS Neuronライブラリのドキュメントを見て進めます。 コンパイルは、普通のCPUのEC2インスタンスやInf1の載ったインスタンスのどちらでも可能です。
次に、公開モデルを例にモデルのコンパイルを行います*1。
import os import torch.neuron from transformers import AutoTokenizer, AutoModelForTokenClassification from transformers import pipeline if __name__ == "__main__": save_dir = "sample_neuron" tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER") model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER") nlp = pipeline("ner", model=model, tokenizer=tokenizer) # Trace example = "My name is Wolfgang and I live in Berlin" max_length = 8 tokens = tokenizer( example, padding="max_length", max_length=max_length, return_tensors="pt", truncation=True, ) example_input = ( tokens["input_ids"], tokens["attention_mask"], tokens["token_type_ids"], ) model_neuron = torch.neuron.trace(model, example_input, strict=False) model.config.update({"traced_sequence_length": max_length}) os.makedirs(save_dir, exist_ok=True) model_neuron.save(os.path.join(save_dir, "neuron_model.pt")) model.config.save_pretrained(save_dir) tokenizer.save_pretrained(save_dir)
save_dir 配下にモデルが出力されていれば、コンパイル完了です。
コンパイル時に気をつけないといけないことは、モデルの入力としてキーワード引数に対応していないことです。 コンパイル時にトレースできるようにタプルで入力を与えます。
多くの場合、モデルはキーワード引数をとるかと思います。 推論時に必要な引数については、第一引数から順に並ぶようにコードを修正しておく必要があります*2。
Inf1のSageMaker推論エンドポイントを構築する
推論コードの用意
推論コードは以下のように準備します。 コンパイル済みのモデルは、固定長の入力を求められます。 コンフィグに追加の設定を保持するようにしています。
class NeuronCompiledBertNERTagger: def __init__(self, model_path: str, compiled_model_path: str): self.model = torch.jit.load(compiled_model_path) self.tokenizer = AutoTokenizer.from_pretrained( model_path, do_lower_case=False ) self.config = AutoConfig.from_pretrained(model_path) def predict(self, sentences: List[str]) -> List[dict]: outputs = [] for sentence in sentences: tokens = self._preprocess(sentence) logits = self._predict(tokens) scores = self._calc_scores(logits) output = self._postprocess(scores, sentence) outputs.append(output) return outputs def _preprocess(self, sentence: str) -> torch.Tensor: return self.tokenizer( sentence, padding="max_length", max_length=self.config.traced_sequence_length., return_tensors="pt", truncation=True, ) def _predict(self, tokens: List[torch.Tensor]) -> np.array: model_input = ( tokens["input_ids"], tokens["attention_mask"], tokens["token_type_ids"], ) logits = self.model(*model_input)["logits"][0] return logits.detach().numpy() def _calc_scores(self, logits: np.array) -> np.array: maxes = np.max(logits, axis=-1, keepdims=True) shifted_exp = np.exp(logits - maxes) scores = shifted_exp / shifted_exp.sum(axis=-1, keepdims=True) return scores def _postprocess(self, scores: np.array, text: str) -> List[dict]: predicted_words = [] label_ids = np.argmax(scores, axis=-1)[1:-1] token_scores = np.max(scores, axis=-1)[1:-1] tokens = self.tokenizer.tokenize(text) for i, (label_id, score, token) in enumerate(zip(label_ids, token_scores, tokens), start=1): predicted_word = { "word": remove_wordpiece_special_character(token), "entity": self.config.id2label[label_id], "index": i, "score": float(score), } predicted_words.append(predicted_word) return predicted_words
コンテナの用意
現在、独自コンテナで推論エンドポイントを構築しています。 Inf1用のHuggingFaceコンテナが用意されているため、これをベースイメージとして独自コンテナを作成します。
以下にDockerfileのサンプルを用意します。
FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-inference-neuron:1.10.2-transformers4.20.1-neuron-py37-sdk1.19.1-ubuntu18.04 # 各種セットアップを行う。 ... ... COPY sagemaker/train.py /opt/ml/code/train.py ENV SAGEMAKER_PROGRAM train.py # ベースイメージは、ENTRYPOINTに以下のスクリプトが設定されている。 # /usr/local/bin/dockerd-entrypoint.py # これだと、任意のserveを走らせられないため、ENTRYPOINTを上書きする。 ENTRYPOINT [""] CMD ["serve"]
コンテナはビルドして、ECRにプッシュしておきます。
モデルの登録
モデルをSageMakerに登録します。 SageMakerコンソールのモデル画面から「モデルの作成」ボタンを押し、必要事項を記入してモデルを作成します。
コンテナ入力オプションとして、「モデルアーティファクトと推論イメージの場所を指定します。」を利用します。
モデルアーティファクトはプッシュしたECRからパスを取得します。
推論イメージの場所は、S3上のs3://bucket/yyyy/xxx.tar.gzというように圧縮済みのモデルを指定します。
あらかじめ、コンパイル済みのモデル一式をtar.gzで圧縮し、所定のS3バケットにPUTしておく必要があります。
環境変数では、「NEURON_RT_NUM_CORES」を設定します。 これは、一プロセスが使用するNEURONコアの数を指定します。 例えば、ml.inf1.xlargeでは4コア利用可能です。4プロセスで起動した際に、1プロセスに1つのInfコア利用する場合、1を設定します。
エンドポイントの作成
エンドポイント設定を作成し、エンドポイントを作成します。 エンドポイント設定作成時に、本番稼働用バリアントに登録したモデルを設定します。 AWSコンソールのバグなのか、一度バリアントを削除してから、再度モデルを登録しないとインスタンスタイプを選ぶことができません。 インスタンスタイプに、ml.inf1.xlargeを選び登録します。
推論してみる
推論は以下のように、立てたエンドポイント名を指定し、boto3クライアントから利用します。
import json import boto3 sagemaker_client = boto3.client("sagemaker-runtime") item = {"text": "My name is Wolf and I live in Germany"} response = sagemaker_client.invoke_endpoint( EndpointName="neuron-ner-endpoint-name", Body=json.dumps(item).encode("utf-8") ) body = json.loads(response.get("Body").read().decode())
速度
リアルタイム推論で、ml.m4.xlarge と ml.inf1.xlarge を比較してみました。
フェアな比較にはなっていませんが、ワーカ数はメモリエラーなどで落ちない範囲で最大にしたところ、以下のようになりました。
- ml.m4.xlarge
- プロセス数：4
- 高速化：線形層の量子化（qint8）
- バッチサイズ：1
- ml.inf1.xlarge
- プロセス数：3
- 高速化：Inf1向けにコンパイル（入力長は512）
- バッチサイズ：1
AWS EC2インスタンス上から、100文をSageMakerの推論エンドポイントを利用して解析した際にかかる時間を測定しました。
ml.m4.xlarge
- 処理時間：21.5秒
- スループット：4.7 件/秒
ml.inf1.xlarge
- 処理時間：3.3秒
- スループット：30 件/秒
単純計算ですが、4.7 件/秒を上回るリクエストが頻繁にくるAPIであれば、Inf1インスタンスにしたほうが台数が少なくリクエストが捌けるようになるため、安くなると思います。
バッチサイズを変えたり、実装を見直したり、コンパイルしたモデルの配置方法を変えたりと、まだまだ早くなる余地はあるかと思います。
お手軽に高速化を行う
コンパイルのお作法はドキュメントを見てしっかり抑える必要があります。 お作法に従えば、手軽にモデルを高速化することができます。 ただし、トレースできるモデルであること、という条件を満たす必要があります。
コンパイル済みのモデルを評価したところスコアや出力はあまり変化していないため、本番に導入予定です。
スループットをあげたい状況やスケールアウトしてコストがかかっている場合には、選択肢の一つとしてInf1インスタンスを利用してみてもいいと思います。
執筆者プロフィール
高橋寛治 Sansan株式会社 技術本部 研究開発部 Data Analysisグループ
阿南工業高等専門学校卒業後に、長岡技術科学大学に編入学。同大学大学院電気電子情報工学専攻修了。在学中は、自然言語処理の研究に取り組み、解析ツールの開発や機械翻訳に関連する研究を行う。大学院を卒業後、2017年にSansan株式会社に入社。キーワード抽出など自然言語処理を生かした研究に取り組む。
▼執筆者による連載記事はこちら
