12 KiB
Export to ONNX
🤗 Transformersモデルを本番環境に展開する際には、モデルを特殊なランタイムおよびハードウェアで読み込み、実行できるように、モデルをシリアライズされた形式にエクスポートすることが必要であるか、その恩恵を受けることができることがあります。
🤗 Optimumは、Transformersの拡張機能であり、PyTorchまたはTensorFlowからモデルをONNXやTFLiteなどのシリアライズされた形式にエクスポートすることを可能にする「exporters」モジュールを提供しています。また、🤗 Optimumは、最大の効率でターゲットハードウェアでモデルをトレーニングおよび実行するためのパフォーマンス最適化ツールも提供しています。
このガイドでは、🤗 Transformersモデルを🤗 Optimumを使用してONNXにエクスポートする方法を示しており、モデルをTFLiteにエクスポートする方法についてはExport to TFLiteページを参照してください。
Export to ONNX
ONNX(Open Neural Network eXchange)は、PyTorchおよびTensorFlowを含むさまざまなフレームワークで深層学習モデルを表現するための共通の一連の演算子とファイル形式を定義するオープンスタンダードです。モデルがONNX形式にエクスポートされると、これらの演算子はニューラルネットワークを介するデータの流れを表す計算グラフ(一般的には「中間表現」と呼ばれる)を構築するために使用されます。
標準化された演算子とデータ型を備えたグラフを公開することで、ONNXはフレームワーク間の切り替えを容易にします。たとえば、PyTorchでトレーニングされたモデルはONNX形式にエクスポートし、それをTensorFlowでインポートすることができます(逆も同様です)。
ONNX形式にエクスポートされたモデルは、以下のように使用できます:
- グラフ最適化や量子化などのテクニックを使用して推論のために最適化。
ORTModelForXXX
クラスを介してONNX Runtimeで実行し、🤗 TransformersでおなじみのAutoModel
APIに従います。- 最適化された推論パイプラインを介して実行し、🤗 Transformersの[
pipeline
]関数と同じAPIを持っています。
🤗 Optimumは、設定オブジェクトを活用してONNXエクスポートをサポートしており、これらの設定オブジェクトは多くのモデルアーキテクチャ用に事前に作成されており、他のアーキテクチャにも簡単に拡張できるように設計されています。
事前に作成された設定のリストについては、🤗 Optimumドキュメントを参照してください。
🤗 TransformersモデルをONNXにエクスポートする方法は2つあります。以下では両方の方法を示します:
- export with 🤗 Optimum via CLI.
- export with 🤗 Optimum with
optimum.onnxruntime
.
Exporting a 🤗 Transformers model to ONNX with CLI
🤗 TransformersモデルをONNXにエクスポートするには、まず追加の依存関係をインストールしてください:
pip install optimum[exporters]
すべての利用可能な引数を確認するには、🤗 Optimumドキュメントを参照してください。または、コマンドラインでヘルプを表示することもできます:
optimum-cli export onnx --help
🤗 Hubからモデルのチェックポイントをエクスポートするには、例えば distilbert/distilbert-base-uncased-distilled-squad
を使いたい場合、以下のコマンドを実行してください:
optimum-cli export onnx --model distilbert/distilbert-base-uncased-distilled-squad distilbert_base_uncased_squad_onnx/
進行状況を示し、結果の model.onnx
が保存される場所を表示するログは、以下のように表示されるはずです:
Validating ONNX model distilbert_base_uncased_squad_onnx/model.onnx...
-[✓] ONNX model output names match reference model (start_logits, end_logits)
- Validating ONNX Model output "start_logits":
-[✓] (2, 16) matches (2, 16)
-[✓] all values close (atol: 0.0001)
- Validating ONNX Model output "end_logits":
-[✓] (2, 16) matches (2, 16)
-[✓] all values close (atol: 0.0001)
The ONNX export succeeded and the exported model was saved at: distilbert_base_uncased_squad_onnx
上記の例は🤗 Hubからのチェックポイントのエクスポートを示しています。ローカルモデルをエクスポートする場合、まずモデルの重みとトークナイザのファイルを同じディレクトリ(local_path
)に保存してください。CLIを使用する場合、🤗 Hubのチェックポイント名の代わりにmodel
引数にlocal_path
を渡し、--task
引数を指定してください。🤗 Optimumドキュメントでサポートされているタスクのリストを確認できます。task
引数が指定されていない場合、タスク固有のヘッドを持たないモデルアーキテクチャがデフォルトで選択されます。
optimum-cli export onnx --model local_path --task question-answering distilbert_base_uncased_squad_onnx/
エクスポートされた model.onnx
ファイルは、ONNX標準をサポートする多くのアクセラレータの1つで実行できます。たとえば、ONNX Runtimeを使用してモデルを読み込み、実行する方法は以下の通りです:
>>> from transformers import AutoTokenizer
>>> from optimum.onnxruntime import ORTModelForQuestionAnswering
>>> tokenizer = AutoTokenizer.from_pretrained("distilbert_base_uncased_squad_onnx")
>>> model = ORTModelForQuestionAnswering.from_pretrained("distilbert_base_uncased_squad_onnx")
>>> inputs = tokenizer("What am I using?", "Using DistilBERT with ONNX Runtime!", return_tensors="pt")
>>> outputs = model(**inputs)
🤗 HubからTensorFlowのチェックポイントをエクスポートするプロセスは、同様です。例えば、Keras organizationから純粋なTensorFlowのチェックポイントをエクスポートする方法は以下の通りです:
optimum-cli export onnx --model keras-io/transformers-qa distilbert_base_cased_squad_onnx/
Exporting a 🤗 Transformers model to ONNX with optimum.onnxruntime
CLIの代わりに、🤗 TransformersモデルをONNXにプログラム的にエクスポートすることもできます。以下のように行います:
>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import AutoTokenizer
>>> model_checkpoint = "distilbert_base_uncased_squad"
>>> save_directory = "onnx/"
>>> # Load a model from transformers and export it to ONNX
>>> ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint, export=True)
>>> tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
>>> # Save the onnx model and tokenizer
>>> ort_model.save_pretrained(save_directory)
>>> tokenizer.save_pretrained(save_directory)
Exporting a model for an unsupported architecture
現在エクスポートできないモデルをサポートするために貢献したい場合、まずoptimum.exporters.onnx
でサポートされているかどうかを確認し、サポートされていない場合は🤗 Optimumに貢献してください。
Exporting a model with transformers.onnx
transformers.onnx
はもはやメンテナンスされていないため、モデルを上記で説明したように🤗 Optimumでエクスポートしてください。このセクションは将来のバージョンで削除されます。
🤗 TransformersモデルをONNXにエクスポートするには、追加の依存関係をインストールしてください:
pip install transformers[onnx]
transformers.onnx
パッケージをPythonモジュールとして使用して、事前に用意された設定を使用してチェックポイントをエクスポートする方法は以下の通りです:
python -m transformers.onnx --model=distilbert/distilbert-base-uncased onnx/
この方法は、--model
引数で定義されたチェックポイントのONNXグラフをエクスポートします。🤗 Hubのいずれかのチェックポイントまたはローカルに保存されたチェックポイントを渡すことができます。エクスポートされたmodel.onnx
ファイルは、ONNX標準をサポートする多くのアクセラレータで実行できます。例えば、ONNX Runtimeを使用してモデルを読み込んで実行する方法は以下の通りです:
>>> from transformers import AutoTokenizer
>>> from onnxruntime import InferenceSession
>>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
>>> session = InferenceSession("onnx/model.onnx")
>>> # ONNX Runtime expects NumPy arrays as input
>>> inputs = tokenizer("Using DistilBERT with ONNX Runtime!", return_tensors="np")
>>> outputs = session.run(output_names=["last_hidden_state"], input_feed=dict(inputs))
必要な出力名(例: ["last_hidden_state"]
)は、各モデルのONNX構成を確認することで取得できます。例えば、DistilBERTの場合、次のようになります:
>>> from transformers.models.distilbert import DistilBertConfig, DistilBertOnnxConfig
>>> config = DistilBertConfig()
>>> onnx_config = DistilBertOnnxConfig(config)
>>> print(list(onnx_config.outputs.keys()))
["last_hidden_state"]
ハブから純粋なTensorFlowのチェックポイントをプログラム的にエクスポートするプロセスは、以下のように同様です:
python -m transformers.onnx --model=keras-io/transformers-qa onnx/
ローカルに保存されたモデルをエクスポートする場合、モデルの重みとトークナイザのファイルを同じディレクトリに保存してください(例: local-pt-checkpoint
)。その後、transformers.onnx
パッケージの --model
引数を希望するディレクトリに向けて設定して、ONNXにエクスポートします:
python -m transformers.onnx --model=local-pt-checkpoint onnx/