Transformer 模型优化工具概述
虽然 ONNX Runtime 在加载 Transformer 模型时会自动应用大多数优化,但一些最新优化尚未集成到 ONNX Runtime 中。可以使用Transformer 优化工具来应用这些额外的优化,以调整模型以获得最佳性能。当 ONNX Runtime 未在加载时应用优化时,此优化工具提供了离线优化 Transformer 模型的能力。
此工具在以下情况下会有所帮助:
- ONNX Runtime 尚未启用特定于 Transformer 的图优化
- 模型可以转换为使用 float16,以在具有 Tensor Cores (如 V100 或 T4) 的 GPU 上通过混合精度提升性能
- 模型输入具有动态轴,这会阻止 ONNX Runtime 由于形状推断而应用某些优化。
- 试验禁用或启用某些融合,以评估对性能或准确性的影响。
用法
支持的模型
有关已使用优化器测试的模型列表,请参阅此页面。
大多数优化需要子图的精确匹配。子图中的任何布局更改都可能导致某些优化无法工作。请注意,不同版本的训练或导出工具可能导致不同的图布局。建议使用最新发布的 PyTorch 和 Transformers 版本。
限制
- 由于 ONNX Runtime 中 Attention 内核的 CUDA 实现,最大注意力头数量为 1024。
- 通常,由于 GPU 内存限制,Longformer 支持的最大序列长度为 4096,其他类型的模型为 1024。
1. 安装 ONNX Runtime
首先,您需要安装 onnxruntime 或 onnxruntime-gpu 包以进行 CPU 或 GPU 推理。要使用 onnxruntime-gpu,需要安装 CUDA 和 cuDNN 并将其 bin 目录添加到 PATH 环境变量。请参阅Python 安装说明。
2. 将 Transformer 模型转换为 ONNX
要将 Transformer 模型转换为 ONNX,请使用torch.onnx或tensorflow-onnx。
GPT-2 模型转换
当使用过去状态时,将 GPT-2 模型从 PyTorch 转换为 ONNX 并不简单。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 操作系统(例如 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
它将在该目录下生成一个 PyTorch 扩展文件,例如“build/lib.linux-x86_64-3.8/longformer_attention.cpython-38-x86_64-linux-gnu.so”。
最后,将 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 模型有三个输入(例如 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。其他三个命令将针对三个引擎中的每一个运行性能测试:OnnxRuntime、PyTorch 和 PyTorch+TorchScript。
如果移除 -o 参数,则优化器脚本不会用于基准测试。
如果您的 GPU(如 V100 或 T4)具有 TensorCore,您可以在上述命令中附加-p fp16
以启用混合精度。在某些仅解码器(例如 GPT2)的生成模型中,您可以在 CUDA EP 上为 SkipLayerNormalization Op 启用严格模式以获得更好的准确性。然而,性能会略有下降。
如果您想在 CPU 上进行基准测试,可以从命令中移除 -g 选项。
请注意,我们目前对 GPT2 和 DistilGPT2 模型的基准测试已禁用输入和输出中的过去状态。
默认情况下,ONNX 模型只有一个输入(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_B1S128
性能分析
profiler.py 可用于对 Transformer 模型进行性能分析。它可以帮助找出模型的瓶颈以及节点或子图上消耗的 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__
基准测试结果可以在这里找到。