使用 ONNX Runtime 实现最佳推理

Optimum 是一个实用程序包，用于使用 ONNX Runtime 等加速运行时构建和运行推理。Optimum 可用于从 Hugging Face Hub 加载优化后的模型，并创建管道以运行加速推理，而无需重写您的 API。

加载

Transformers 模型

将您的模型导出到 ONNX 格式后，您可以通过将 AutoModelForXxx 类替换为相应的 ORTModelForXxx 来加载它。

  from transformers import AutoTokenizer, pipeline
- from transformers import AutoModelForQuestionAnswering
+ from optimum.onnxruntime import ORTModelForQuestionAnswering

- model = AutoModelForQuestionAnswering.from_pretrained("meta-llama/Llama-3.2-1B) # PyTorch checkpoint
+ model = ORTModelForQuestionAnswering.from_pretrained("onnx-community/Llama-3.2-1B", subfolder="onnx") # ONNX checkpoint
  tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")

  pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
  result = pipe("He never went out without a book under his arm")

有关我们支持的所有 ORTModelForXxx 的更多信息，请参阅我们的文档

Diffusers 模型

将您的模型导出到 ONNX 格式后，您可以通过将 DiffusionPipeline 类替换为相应的 ORTDiffusionPipeline 来加载它。

- from diffusers import DiffusionPipeline
+ from optimum.onnxruntime import ORTDiffusionPipeline

  model_id = "runwayml/stable-diffusion-v1-5"
- pipeline = DiffusionPipeline.from_pretrained(model_id)
+ pipeline = ORTDiffusionPipeline.from_pretrained(model_id, revision="onnx")
  prompt = "sailing ship in storm by Leonardo da Vinci"
  image = pipeline(prompt).images[0]

动态将您的模型转换为 ONNX

如果您的模型尚未转换为 ONNX，ORTModel 包含一个方法，可以动态将您的模型转换为 ONNX。只需在 from_pretrained() 方法中传递 export=True，您的模型将被动态加载并转换为 ONNX。

>>> from optimum.onnxruntime import ORTModelForSequenceClassification

>>> # Load the model from the hub and export it to the ONNX format
>>> model_id = "distilbert-base-uncased-finetuned-sst-2-english"
>>> model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)

将您的模型推送到 Hub

您也可以直接在您的模型上调用 push_to_hub，以将其上传到 Hub。

>>> from optimum.onnxruntime import ORTModelForSequenceClassification

>>> # Load the model from the hub and export it to the ONNX format
>>> model_id = "distilbert-base-uncased-finetuned-sst-2-english"
>>> model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)

>>> # Save the converted model locally
>>> output_dir = "a_local_path_for_convert_onnx_model"
>>> model.save_pretrained(output_dir)

# Push the onnx model to HF Hub
>>> model.push_to_hub(output_dir, repository_id="my-onnx-repo")

< > 更新在 GitHub 上