これはTorchScriptを使用した実験の最初であり、可変入力サイズのモデルに対するその能力をまだ探求中です。これは私たちの関心の焦点であり、今後のリリースでは、より柔軟な実装や、PythonベースのコードとコンパイルされたTorchScriptを比較するベンチマークを含む、より多くのコード例で詳細な分析を行います。
TorchScriptのドキュメントによれば:
TorchScriptは、PyTorchコードから直列化および最適化可能なモデルを作成する方法です。
TorchScriptを使用すると、効率志向のC++プログラムなど、他のプログラムでモデルを再利用できるようになります。PyTorchベースのPythonプログラム以外の環境で🤗 Transformersモデルをエクスポートして使用するためのインターフェースを提供しています。ここでは、TorchScriptを使用してモデルをエクスポートし、使用する方法を説明します。
モデルをエクスポートするには、次の2つの要件があります:
torchscript
フラグを使用したモデルのインスタンス化- ダミーの入力を使用したフォワードパス
これらの必要条件は、以下で詳細に説明されているように、開発者が注意する必要があるいくつかのことを意味します。
torchscript
フラグは、ほとんどの🤗 Transformers言語モデルにおいて、Embedding
レイヤーとDecoding
レイヤー間で重みが連結されているため必要です。
TorchScriptでは、重みが連結されているモデルをエクスポートすることはできませんので、事前に重みを切り離して複製する必要があります。
torchscript
フラグを使用してインスタンス化されたモデルは、Embedding
レイヤーとDecoding
レイヤーが分離されており、そのため後でトレーニングしてはいけません。
トレーニングは、これらの2つのレイヤーを非同期にする可能性があり、予期しない結果をもたらす可能性があります。
言語モデルヘッドを持たないモデルには言及しませんが、これらのモデルには連結された重みが存在しないため、torchscript
フラグなしで安全にエクスポートできます。
ダミー入力はモデルのフォワードパスに使用されます。入力の値はレイヤーを通じて伝播される間、PyTorchは各テンソルに実行された異なる操作を追跡します。これらの記録された操作は、モデルのトレースを作成するために使用されます。
トレースは入力の寸法に対して作成されます。そのため、ダミー入力の寸法に制約され、他のシーケンス長やバッチサイズでは動作しません。異なるサイズで試すと、以下のエラーが発生します:
`The expanded size of the tensor (3) must match the existing size (7) at non-singleton dimension 2`
お勧めしますのは、モデルの推論中に供給される最大の入力と同じ大きさのダミー入力サイズでモデルをトレースすることです。パディングを使用して不足値を補完することもできます。ただし、モデルがより大きな入力サイズでトレースされるため、行列の寸法も大きくなり、より多くの計算が発生します。
異なるシーケンス長のモデルをエクスポートする際に、各入力に対して実行される演算の総数に注意して、パフォーマンスを密接にフォローすることをお勧めします。
このセクションでは、モデルの保存と読み込み、および推論にトレースを使用する方法を示します。
TorchScriptでBertModel
をエクスポートするには、BertConfig
クラスからBertModel
をインスタンス化し、それをファイル名traced_bert.pt
でディスクに保存します:
from transformers import BertModel, BertTokenizer, BertConfig
import torch
enc = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
# Tokenizing input text
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)
# Masking one of the input tokens
masked_index = 8
tokenized_text[masked_index] = "[MASK]"
indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
# Creating a dummy input
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]
# Initializing the model with the torchscript flag
# Flag set to True even though it is not necessary as this model does not have an LM Head.
config = BertConfig(
vocab_size_or_config_json_file=32000,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
torchscript=True,
)
# Instantiating the model
model = BertModel(config)
# The model needs to be in evaluation mode
model.eval()
# If you are instantiating the model with *from_pretrained* you can also easily set the TorchScript flag
model = BertModel.from_pretrained("google-bert/bert-base-uncased", torchscript=True)
# Creating the trace
traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "traced_bert.pt")
以前に保存した BertModel
、traced_bert.pt
をディスクから読み込んで、以前に初期化した dummy_input
で使用できます。
loaded_model = torch.jit.load("traced_bert.pt")
loaded_model.eval()
all_encoder_layers, pooled_output = loaded_model(*dummy_input)
トレースモデルを使用して推論を行うには、その __call__
ダンダーメソッドを使用します。
traced_model(tokens_tensor, segments_tensors)
AWSはクラウドでの低コストで高性能な機械学習推論向けに Amazon EC2 Inf1 インスタンスファミリーを導入しました。Inf1インスタンスはAWS Inferentiaチップによって駆動され、ディープラーニング推論ワークロードに特化したカスタムビルドのハードウェアアクセラレータです。AWS Neuron はInferentia用のSDKで、トランスフォーマーモデルをトレースして最適化し、Inf1に展開するためのサポートを提供します。
Neuron SDK が提供するもの:
- クラウドでの推論のためにTorchScriptモデルをトレースして最適化するための、1行のコード変更で使用できる簡単なAPI。
- 改善されたコストパフォーマンス のためのボックス外のパフォーマンス最適化。
- PyTorch または TensorFlow で構築されたHugging Faceトランスフォーマーモデルへのサポート。
BERT(Bidirectional Encoder Representations from Transformers)アーキテクチャやその変種(distilBERT や roBERTa など)に基づくトランスフォーマーモデルは、非生成タスク(抽出型質問応答、シーケンス分類、トークン分類など)において、Inf1上で最適に動作します。ただし、テキスト生成タスクも AWS Neuron MarianMT チュートリアル に従ってInf1上で実行できます。Inferentiaでボックス外で変換できるモデルに関する詳細情報は、Neuronドキュメンテーションの Model Architecture Fit セクションにあります。
モデルをAWS Neuronに変換するには、Neuron SDK 環境 が必要で、AWS Deep Learning AMI に事前に構成されています。
モデルをAWS NEURON用に変換するには、PythonでTorchScriptを使用する と同じコードを使用して BertModel
をトレースします。Python APIを介してNeuron SDKのコンポーネントにアクセスするために、torch.neuron
フレームワーク拡張をインポートします。
from transformers import BertModel, BertTokenizer, BertConfig
import torch
import torch.neuron
次の行を変更するだけで済みます。
- torch.jit.trace(model, [tokens_tensor, segments_tensors])
+ torch.neuron.trace(model, [token_tensor, segments_tensors])
これにより、Neuron SDKはモデルをトレースし、Inf1インスタンス向けに最適化します。
AWS Neuron SDKの機能、ツール、サンプルチュートリアル、最新のアップデートについて詳しく知りたい場合は、AWS NeuronSDK ドキュメンテーション をご覧ください。