Optimum 1.2 推理发布 添加 ONNX Runtime 加速管道

TL;DR

Optimum 1.2 引入了对 Hugging Face Transformers 管道的 ONNX Runtime 推理支持，允许用户将标准 AutoModel 类替换为 ORTModel 等价类，应用图优化和动态量化，并在保持原始精度 >99% 的同时实现最高 2 倍的延迟降低。

Optimum 是什么？

Optimum 是一个开源库，扩展了 Hugging Face Transformers，提供统一的 API 用于在加速硬件上进行性能优化。它提供加速训练、量化、图优化以及现在的推理工具，支持诸如 Graphcore IPU、Habana Gaudi 和 ONNX Runtime 等硬件。

Optimum 1.2 中的新推理和管道特性

  • API 兼容性 – ORTModelForXxx 类是 AutoModelForXxx 的直接替代。示例：
    from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForQuestionAnswering

model = ORTModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2")
tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
  • ONNX Runtime 集成 – 模型被导出为 ONNX 并使用 ONNX Runtime 执行，获得算子融合、常量折叠以及硬件特定内核的优势。
  • 优化与量化 – ORTOptimizer 执行图级优化；ORTQuantizer 进行动态量化（例如 AVX‑512 VNNI），以缩小模型体积并提升延迟。
  • Hub 兼容性 – 优化后的检查点可以推送到 Hugging Face Hub 并从中拉取，实现加速模型的社区共享。
  • 管道安全层 – optimum.pipelines.pipeline 验证任务与模型的兼容性，并在不支持的配置下抛出错误。

端到端教程：加速 RoBERTa 的问答任务

博客文章在 AWS m5.xlarge 实例上演示了一个六步工作流：

  1. 安装 Optimum 与 ONNX Runtime
    pip install "optimum[onnxruntime]==1.2.0"
  2. 将 Transformers 模型转换为 ONNX
    from optimum.onnxruntime import ORTModelForQuestionAnswering
model = ORTModelForQuestionAnswering.from_pretrained(
    "deepset/roberta-base-squad2", from_transformers=True)
model.save_pretrained("onnx")
  3. 应用图优化
    from optimum.onnxruntime import ORTOptimizer, OptimizationConfig
optimizer = ORTOptimizer.from_pretrained("deepset/roberta-base-squad2", feature="question-answering")
optimizer.export(
   
   
    optimization_config=OptimizationConfig(optimization_level=99),
)
  4. 动态量化
    from optimum.onnxruntime import ORTQuantizer, AutoQuantizationConfig
quantizer = ORTQuantizer.from_pretrained("deepset/roberta-base-squad2", feature="question-answering")
qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=True)
quantizer.export(
   
   
    quantization_config=qconfig,
)
  5. 使用 Transformers Pipelines 进行推理
    from transformers import pipeline, AutoTokenizer
from optimum.onnxruntime import ORTModelForQuestionAnswering
tokenizer = AutoTokenizer.from_pretrained("onnx")
model = ORTModelForQuestionAnswering.from_pretrained("onnx", file_name="model-quantized.onnx")
qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
  6. 评估精度和延迟
    • 在 SQuAD‑v2 上的精度：原始 82.15 F1，优化后 82.15 F1，量化后 81.83 F1（≈原始的 99.6 %）。
    • 在 2 核 CPU 上的延迟：原始 ≈ 117 ms，优化 + 量化 ≈ 65 ms（≈2 倍加速）。

当前限制

  • 模型大小 – 只能从 Hub 加载小于 2 GB 的模型。
  • Seq2Seq 支持 – 如摘要（例如 T5）等任务尚未提供。
  • 过去的键值 – 因果语言模型（例如 GPT‑2）尚未使用缓存的注意力状态。
  • 本地缓存 – 优化后的 ONNX 文件在下载后不会本地缓存。

常见问题

  • 支持的任务 – 特征提取、文本分类、标记分类、问答、零样本分类和文本生成。
  • 支持的模型 – 任何可通过 transformers.onnx 导出的模型，包括 BERT、ALBERT、RoBERTa、XLM‑R、DistilBERT、GPT‑2 等。
  • 支持的运行时 – 当前为 ONNX Runtime；计划支持其他运行时（TensorRT、AWS‑Neuron）。
  • GPU 使用 – 安装 optimum[onnxruntime-gpu] 可自动启用 GPU 提供者。
  • 加载优化模型 – 使用 ORTModelForXXX.from_pretrained(..., file_name="model‑optimized.onnx") 方式。

路线图与未来工作

Optimum 旨在成为 transformer 加速的参考工具包。即将到来的里程碑包括：

  • 添加语音（Wav2Vec 2.0）和视觉（ViT）模型支持。
  • 集成 OrtValue 和 IOBinding 以获得更低层次的性能提升。
  • 提供更简便的加速模型评估工具。
  • 将运行时支持扩展到 TensorRT、AWS‑Neuron 等提供者。

如果您有兴趣贡献或了解更多，Optimum 仓库、Hugging Face 论坛以及作者的社交渠道均可进行讨论。

Sources