Transformersオプティマイザ

Transformerモデル最適化ツールの概要

ONNX Runtimeは、トランスフォーマーモデルをロードする際にほとんどの最適化を自動的に適用しますが、ONNX Runtimeにまだ統合されていない最新の最適化もいくつかあります。これらの追加の最適化は、トランスフォーマー最適化ツールを使用して適用し、最高のパフォーマンスを得るためにモデルを調整できます。この最適化ツールは、ONNX Runtimeがロード時に最適化を適用しないシナリオでトランスフォーマーモデルを最適化するオフライン機能を提供します。

このツールは、次の場合に役立ちます。

ONNX Runtimeにまだトランスフォーマー固有のグラフ最適化が有効になっていない場合
Tensorコアを搭載したGPU（V100やT4など）で混合精度を使用してパフォーマンスを向上させるために、モデルをfloat16を使用するように変換できる場合
モデルに入力に動的軸があり、形状推論のためにONNX Runtimeによって一部の最適化が適用されない場合
パフォーマンスや精度への影響を評価するために、一部の融合を無効にしたり有効にしたりして実験する場合

使用法：

ONNX Runtimeのインストール
トランスフォーマーモデルをONNXに変換
モデルオプティマイザツールの実行
モデルのベンチマークとプロファイリング

サポートされているモデル

オプティマイザでテストされたモデルのリストについては、このページを参照してください。

ほとんどの最適化には、サブグラフの完全一致が必要です。サブグラフのレイアウトが変更されると、一部の最適化が機能しない場合があります。トレーニングツールやエクスポートツールのバージョンが異なると、グラフのレイアウトが異なる場合があることに注意してください。PyTorchとTransformersの最新のリリースバージョンを使用することをお勧めします。

制限事項

ONNX RuntimeのAttentionカーネルのCUDA実装により、アテンションヘッドの最大数は1024です。
通常、GPUメモリの制約により、Longformerの最大サポートシーケンス長は4096、その他の種類のモデルでは1024です。

1. ONNX Runtimeのインストール

まず、CPUまたはGPU推論用にonnxruntimeまたはonnxruntime-gpuパッケージをインストールする必要があります。onnxruntime-gpuを使用するには、CUDAとcuDNNをインストールし、それらのbinディレクトリをPATH環境変数に追加する必要があります。Pythonのインストール手順を参照してください。

2. トランスフォーマーモデルをONNXに変換する

トランスフォーマーモデルをONNXに変換するには、torch.onnxまたはtensorflow-onnxを使用します。

Huggingface transformersには、事前学習済みモデルをONNXにエクスポートする例を示すノートブックがあります。
tf2onnxについては、このBERTチュートリアルを参照してください。

GPT-2モデルの変換

過去の状態が使用されている場合、PyTorchからONNXへのGPT-2モデルの変換は簡単ではありません。ツールconvert_to_onnxが役立ちます。

次のようなコマンドを使用して、事前学習済みのPyTorch GPT-2モデルを特定の精度（float32、float16）のONNXに変換できます。

python -m onnxruntime.transformers.models.gpt2.convert_to_onnx -m gpt2 --model_class GPT2LMHeadModel --output gpt2.onnx -p fp32

python -m onnxruntime.transformers.models.gpt2.convert_to_onnx -m distilgpt2 --model_class GPT2LMHeadModel --output distilgpt2.onnx -p fp16 --use_gpu --optimize_onnx --auto_mixed_precision

このツールは、同じランダム入力が与えられた場合に、ONNXモデルと対応するPyTorchモデルが同じ出力を生成するかどうかも検証します。

Longformerモデルの変換

要件：Linux OS（例：Ubuntu 18.04または20.04）および次のようなPyTorch 1.9.*のPython環境：

conda create -n longformer python=3.8
conda activate longformer
pip install torch==1.9.1+cpu torchvision==0.10.1+cpu torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install onnx transformers==4.18.0 onnxruntime numpy

次に、torch拡張機能のソースをビルドします。

cd onnxruntime/python/tools/transformers/models/longformer/torch_extensions
python setup.py install

これにより、ディレクトリ内に「build/lib.linux-x86_64-3.8/longformer_attention.cpython-38-x86_64-linux-gnu.so」のようなPyTorch拡張ファイルが生成されます。

最後に、longformerモデルを次のようにONNXモデルに変換します。

cd ..
python convert_to_onnx.py -m longformer-base-4096

エクスポートされたONNXモデルは、現在GPUでのみ実行できます。

3. モデルオプティマイザツールを実行する

すべてのオプティマイザオプションについては、Githubを参照してください。

Pythonコードでは、次のようにオプティマイザを使用できます。

from onnxruntime.transformers import optimizer
optimized_model = optimizer.optimize_model("bert.onnx", model_type='bert', num_heads=12, hidden_size=768)
optimized_model.convert_float_to_float16()
optimized_model.save_model_to_file("bert_fp16.onnx")

コマンドラインを使用することもできます。BERT-largeモデルを混合精度（float16）を使用するように最適化する例：

python -m onnxruntime.transformers.optimizer --input bert_large.onnx --output bert_large_fp16.onnx --num_heads 16 --hidden_size 1024 --float16

こちらから最新のスクリプトファイルをダウンロードすることもできます。次に、次のように実行します。

python optimizer.py --input bert.onnx --output bert_opt.onnx --model_type bert

BERTモデルの検証

BERTモデルに3つの入力（input_ids、token_type_ids、attention_maskなど）がある場合、スクリプトcompare_bert_results.pyを使用して簡単な検証を行うことができます。このツールは、いくつかの偽の入力データを生成し、元のモデルと最適化されたモデルの両方の結果を比較します。出力がすべて近い場合は、最適化されたモデルを安全に使用できます。

CPU用に最適化されたモデルを検証する例：

python -m onnxruntime.transformers.compare_bert_results --baseline_model original_model.onnx --optimized_model optimized_model_cpu.onnx --batch_size 1 --sequence_length 128 --samples 100

GPUの場合は、コマンドに—use_gpuを追加してください。

4. モデルのベンチマークとプロファイリング

ベンチマーク

bashスクリプトrun_benchmark.shを使用してベンチマークを実行できます。実行する前に、bashスクリプトを変更してオプション（モデル、バッチサイズ、シーケンス長、ターゲットデバイスなど）を選択できます。

このbashスクリプトは、benchmark.pyスクリプトを呼び出して、Huggingface Transformersの事前学習済みモデルに対するOnnxRuntime、PyTorch、またはPyTorch+TorchScriptの推論パフォーマンスを測定します。

Benchmark.py

run_benchmark.shを使用する場合は、benchmark.pyを直接使用する必要はありません。詳細を知りたくない場合は、このセクションをスキップできます。

以下は、GPU上の事前学習済みモデルbert-base-casedでbenchmark.pyを実行する例です。

python -m onnxruntime.transformers.benchmark -g -m bert-base-cased -o -v -b 0
python -m onnxruntime.transformers.benchmark -g -m bert-base-cased -o
python -m onnxruntime.transformers.benchmark -g -m bert-base-cased -e torch
python -m onnxruntime.transformers.benchmark -g -m bert-base-cased -e torchscript

最初のコマンドは、（最適化の前後の両方の）ONNXモデルを生成しますが、バッチサイズが0であるためパフォーマンステストは実行しません。他の3つのコマンドは、3つのエンジン（OnnxRuntime、PyTorch、PyTorch+TorchScript）のそれぞれでパフォーマンステストを実行します。

-oパラメータを削除すると、ベンチマークでオプティマイザスクリプトは使用されません。

GPU（V100やT4など）にTensorCoreがある場合は、上記のコマンドに-p fp16を追加して混合精度を有効にすることができます。一部のデコーダーのみ（GPT2など）の生成モデルでは、CUDA EPでSkipLayerNormalization Opの厳密モードを有効にして、より良い精度を達成できます。ただし、パフォーマンスは少し低下します。

CPUでベンチマークを実行したい場合は、コマンドから-gオプションを削除できます。

GPT2およびDistilGPT2モデルの現在のベンチマークでは、入力と出力から過去の状態が無効になっていることに注意してください。

デフォルトでは、ONNXモデルには1つの入力（input_ids）しかありません。-iパラメータを使用して、複数の入力を持つモデルをテストできます。たとえば、コマンドラインに「-i 3」を追加して、3つの入力（input_ids、token_type_ids、attention_mask）を持つbertモデルをテストできます。このオプションは現在OnnxRuntimeのみをサポートしています。

パフォーマンステスト

bert_perf_test.pyを使用して、BERTモデルの推論パフォーマンスを確認できます。以下に例を示します。

python -m onnxruntime.transformers.bert_perf_test --model optimized_model_cpu.onnx --batch_size 1 --sequence_length 128

GPUの場合は、コマンドに—use_gpuを追加してください。

テストが終了すると、モデルディレクトリにperf_results_CPU_B1_S128_<date_time>.txtやperf_results_GPU_B1_S128_<date_time>.txtのようなファイルが出力されます。

プロファイリング

profiler.pyを使用して、トランスフォーマーモデルのプロファイリングを実行できます。これにより、モデルのボトルネックや、ノードまたはサブグラフで費やされたCPU時間を特定できます。

コマンドの例：

python -m onnxruntime.transformers.profiler --model bert.onnx --batch_size 8 --sequence_length 128 --samples 1000 --dummy_inputs bert --thread_num 8 --kernel_time_only
python -m onnxruntime.transformers.profiler --model gpt2.onnx --batch_size 1 --sequence_length 1 --past_sequence_length 128 --samples 1000 --dummy_inputs gpt2 --use_gpu
python -m onnxruntime.transformers.profiler --model longformer.onnx --batch_size 1 --sequence_length 4096 --global_length 8 --samples 1000 --dummy_inputs longformer --use_gpu

onnxruntime_profile__<date_time>.jsonのような結果ファイルが現在のディレクトリに出力されます。ノードの概要、最も高価なノード、および演算子タイプでグループ化された結果がコンソールに出力されます。

ベンチマークの結果はこちらで確認できます。