Transformer 模型优化工具概述

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

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

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

用法

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

支持的模型

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

大多数优化需要子图的精确匹配。子图中的任何布局更改都可能导致某些优化无法工作。请注意，不同版本的训练或导出工具可能导致不同的图布局。建议使用最新发布的 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。

  • Huggingface Transformers 有一个notebook，其中展示了将预训练模型导出到 ONNX 的示例。

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

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.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 将输出到当前目录。节点的摘要、开销最大的节点以及按操作符类型分组的结果将打印到控制台。

基准测试结果可以在这里找到。