Transformer 模型优化工具概览

虽然 ONNX Runtime 在加载 transformer 模型时会自动应用大多数优化，但一些最新的优化尚未集成到 ONNX Runtime 中。这些额外的优化可以使用 transformer 优化工具来应用，以调整模型以获得最佳性能。该优化工具提供了一种离线优化 transformer 模型的能力，适用于 ONNX Runtime 在加载时未应用优化的情况。

此工具在以下情况下很有帮助：

  • ONNX Runtime 尚未启用 transformer 特定的图优化
  • 模型可以转换为使用 float16，以便在具有 Tensor Core 的 GPU（如 V100 或 T4）上使用混合精度提高性能
  • 模型的输入具有动态轴，由于形状推断，这会阻止 ONNX Runtime 应用某些优化。
  • 尝试禁用或启用某些融合以评估对性能或准确性的影响。

用法

  1. 安装 ONNX Runtime
  2. 将 transformer 模型转换为 ONNX
  3. 运行模型优化工具
  4. 对模型进行基准测试和性能分析

支持的模型

有关已使用优化器测试过的模型列表，请参阅此页面。

大多数优化需要子图的精确匹配。子图中的任何布局更改都可能导致某些优化失效。请注意，不同版本的训练或导出工具可能导致不同的图布局。建议使用最新发布的 PyTorch 和 Transformers 版本。

限制

  • 由于 ONNX Runtime 中 Attention 内核的 CUDA 实现，最大 attention head 数量为 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。

  • Huggingface transformers 提供了一个 notebook 演示了导出预训练模型到 ONNX 的示例。

  • 对于 tf2onnx，请参考此 BERT 教程。

GPT-2 模型转换

当使用 past state 时，将 GPT-2 模型从 PyTorch 转换为 ONNX 并不直接。工具 convert_to_onnx 可以提供帮助。

你可以使用以下命令将预训练的 PyTorch GPT-2 模型转换为 ONNX，指定精度（float32，float16）

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) 和一个 Python 环境，其中包含 PyTorch 1.9.*，如下所示：

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 extensions 的源代码

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 以启用混合精度。在一些基于 decoder-only（如 GPT2）的生成模型中，你可以在 CUDA EP 上为 SkipLayerNormalization Op 启用严格模式以获得更好的准确性。但是，性能会略有下降。

如果你想在 CPU 上进行基准测试，可以移除命令中的 -g 选项。

请注意，我们当前在 GPT2 和 DistilGPT2 模型上的基准测试已禁用输入和输出中的 past state。

默认情况下，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.txt 或 perf_results_GPU_B1_S128_.txt。

性能分析

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__.json 将输出到当前目录。节点摘要、最耗时节点和按算子类型分组的结果将打印到控制台。

基准测试结果可在此处找到。